[Fix] Update type hint for position_ids, update max_lengths type in LengthGroupedSampler, and handle fp32 lmhead hf weight (#1636)

* [Fix] Update type hint for position_ids in MultimodalTrainInfo and assert max_lengths type in LengthGroupedSampler

- Changed position_ids to allow None in MultimodalTrainInfo.
- Updated max_lengths assertion to include Column type in LengthGroupedSampler.
- Adjusted reward assignment in JudgerController to use "score" instead of "weighted_score".
- Enhanced tensor handling in BaseModel for FSDP with specific conditions for lm_head.weight.

* fix lint

* fix ut

* fix ut

Author: RangiLyu

tests/ray/test_judger.py

@@ -189,7 +189,7 @@ def test_gsm8k_multi_judger(self):
         pg = AutoCPUWorkers.build_placement_group(cpu_resources_config)
         judger_controller = JudgerController.remote(judger_cfg, pg)
         res3 = ray.get(judger_controller.run.remote(FAKE_JUDGER_INPUT_ITEM_MULTI_SOURCE))
-        self.assertEqual(res3.reward["weighted_score"], 1.0)  # weighted_score为固定字段，表示加权后的reward
+        self.assertEqual(res3.reward["score"], 1.0)
 
     def test_gsm8k_judger_score(self):
         """Test the judger functionality with single and multiple data sources."""

xtuner/v1/data_proto/rl_data.py

@@ -86,7 +86,7 @@ class RLUIDItem(BaseModel):
 class MultimodalTrainInfo(TypedDict):
     pixel_values: NotRequired[torch.Tensor | RayObjectRef | None]  # type: ignore[valid-type]
     image_grid_thw: NotRequired[torch.Tensor]
-    position_ids: NotRequired[torch.Tensor]
+    position_ids: NotRequired[torch.Tensor | None]
 
 
 class RLDatasetItem(BaseModel):
@@ -206,7 +206,7 @@ class RLJudgerResponseItem(BaseModel):
 
     Attributes:
         uid (Optional[int]): A unique ID to identify which input the result corresponds to.
-        reward (Dict[str, Any]): A dictionary of reward scores, e.g., {"judger_type": reward_score, "weighted_scores": score}.
+        reward (Dict[str, Any]): A dictionary of reward scores, e.g., {"score": score}.
         extra_info (Dict[str, Any]): Additional user-defined information.
     """
 

xtuner/v1/datasets/sampler.py

@@ -15,6 +15,11 @@
 from .packing import _LegacySoftPackDataset
 
 
+try:
+    from datasets.arrow_dataset import Column
+except ImportError:
+    Column = type("Column", (), {})  # Dummy type for backward compatibility
+
 logger = get_logger()
 
 
@@ -221,7 +226,7 @@ def __init__(
         self.group_size = self.world_size
 
         self.max_lengths = self.dataset.longest
-        assert isinstance(self.max_lengths, (list, tuple))
+        assert isinstance(self.max_lengths, (list, tuple, Column))
 
         self.global_batch_size = global_batch_size
 

xtuner/v1/model/base.py

@@ -1017,8 +1017,24 @@ def _get_same_hf_param(
                 buffer_name_list.append(load_spec.hf_keys[0])
                 continue
             local_tensor = param._local_tensor if isinstance(param, DTensor) else param
-            local_tensor = local_tensor.to(dtype=self._get_save_dtype(load_spec.hf_keys[0], torch.bfloat16))
-            tensor_size = self._get_tensor_size(param, dtype)
+            if (
+                self.fsdp_config is not None
+                and self.fsdp_config.fp32_lm_head
+                and load_spec.hf_keys[0] == "lm_head.weight"
+            ):
+                logger.info(f"handling same hf param: {load_spec.hf_keys} separately")
+                lm_head_tensor_list = self._fsdp_foreach_allgather([local_tensor], [load_spec])
+                lm_head_tensor_list = [
+                    self.param_to_safetensor(safetensor, name)
+                    for safetensor, name in zip(lm_head_tensor_list, load_spec.hf_keys.copy())
+                ]
+                lm_head_tensor_list = [t.to(device=device) for t in lm_head_tensor_list]
+                yield load_spec.hf_keys.copy(), lm_head_tensor_list
+                del lm_head_tensor_list, local_tensor
+                continue
+            else:
+                local_tensor = local_tensor.to(dtype=self._get_save_dtype(load_spec.hf_keys[0], torch.bfloat16))
+                tensor_size = self._get_tensor_size(param, dtype)
             if safetensor_size + tensor_size > bucket_size and tensor_list:
                 if self.fsdp_mesh is not None:
                     gathered_tensor_list = self._fsdp_foreach_allgather(tensor_list, load_spec_list)

xtuner/v1/ray/judger/controller.py

@@ -260,7 +260,7 @@ async def run(
                 for name, weight in data_source.items():
                     if name in item.reward:
                         final_reward += item.reward[name]["score"] * weight
-                item.reward["weighted_score"] = final_reward
+                item.reward["score"] = final_reward
         else:
             judger_response_item = await self._call_custom_reward_judger(
                 self.reward_judger, self.reward_judger_names, group_data_item