add ut and unify training input to WorkerLogItem

Author: YanhuiDua

tests/ray/test_grpo_train.py

@@ -55,7 +55,6 @@ def setUp(self):
             rewards = [item['reward'] for item in group]
             rewards = torch.tensor(rewards, dtype=torch.float32)
             advantages = (rewards - rewards.mean(0)) / (rewards.std(0) + 1e-8)
-
             for i in range(self.prompt_repeat_k):
                 item = group[i]
                 response_ids = tokenizer(item['response'], return_tensors='pt')['input_ids'].flatten().tolist()
@@ -67,7 +66,7 @@ def setUp(self):
                     dict(
                         seq_ctx=SequenceContext.from_input_ids((input_ids, ), device="cpu"),
                         shifted_labels=shifted_labels,
-                        advantage=advantages[i].item(),
+                        advantages=advantages[i],
                     )
                 )
         self.data_batches = data_batches
@@ -126,8 +125,125 @@ def build_train_controller(self):
         ray.get(train_controller.__ray_ready__.remote())
         return train_controller
 
-    def test_grpo_train_and_save(self):
+    # def test_grpo_train_and_save(self):
+    #     train_controller = self.build_train_controller()
+    #     ray.get(train_controller.fit.remote(self.data_batches, pack_max_length=8192, rollout_idx=0))
+    #     save_path = os.path.join(self.temp_dir, "hf_test")
+    #     ray.get(train_controller.save_hf.remote(str(save_path)))
+
+    def _create_dummy_item(self, length: int):
+        """Helper to create a dummy WorkerInputItem"""
+        input_ids = torch.ones(1, length, dtype=torch.long)
+        cu_seq_lens_q = torch.tensor([0, length], dtype=torch.int32)
+        cu_seq_lens_k = torch.tensor([0, length], dtype=torch.int32)
+        max_length_q = torch.tensor(length, dtype=torch.int32)
+        max_length_k = torch.tensor(length, dtype=torch.int32)
+        seq_ctx = SequenceContext(
+            input_ids=input_ids,
+            cu_seq_lens_q=cu_seq_lens_q,
+            cu_seq_lens_k=cu_seq_lens_k,
+            max_length_q=max_length_q,
+            max_length_k=max_length_k,
+            num_padding=0,
+            device="cpu",
+        )
+        return {
+            "seq_ctx": seq_ctx,
+            "shifted_labels": torch.ones(1, length, dtype=torch.long),
+            "advantages": torch.rand(1, 1, dtype=torch.float),
+            "rollout_logprobs": torch.ones(1, length, dtype=torch.float),
+        }
+
+    def test_controller_logic(self):
+        """
+        Unit tests for RawTrainingController internal logic using the real Ray actor:
+        - _balance_split_batch
+        - _create_padding_item
+        - _rearrange_batch_for_pack
+        - _pad_and_pack_batches
+        """
+        # 1. Build the real train controller
         train_controller = self.build_train_controller()
-        ray.get(train_controller.fit.remote(self.data_batches, pack_max_length=1024, rollout_idx=0))
-        save_path = os.path.join(self.temp_dir, "hf_test")
-        ray.get(train_controller.save_hf.remote(str(save_path)))
+        pack_max_length = 100
+
+        # --- Test 1: _balance_split_batch ---
+        print("Testing _balance_split_batch...")
+        # Input: 4 items with lengths 10, 20, 30, 40
+        items = [self._create_dummy_item(l) for l in [10, 20, 30, 40]]
+        dp_size = 2
+        
+        # Call remote method
+        # 10, 20, 30, 40 -> sum 100 -> avg 50.
+        # Expected split: [10, 40] (sum 50) and [20, 30] (sum 50)
+        result = ray.get(train_controller._balance_split_batch.remote(items, dp_size))
+        
+        self.assertEqual(len(result), 2)
+        self.assertEqual(len(result[0]), 2)
+        self.assertEqual(len(result[1]), 2)
+        
+        # Verify balance
+        len_group0 = sum(item["seq_ctx"].input_ids.shape[1] for item in result[0])
+        len_group1 = sum(item["seq_ctx"].input_ids.shape[1] for item in result[1])
+        self.assertEqual(len_group0, 50)
+        self.assertEqual(len_group1, 50)
+
+        # --- Test 2: _rearrange_batch_for_pack ---
+        print("Testing _rearrange_batch_for_pack...")
+        # Input: [40, 40, 30], max=100. With get_seqlen_balanced_partitions, it should be packed as [40, 30] and [40]
+        items_pack = [self._create_dummy_item(l) for l in [40, 40, 30]]
+        batches = ray.get(train_controller._rearrange_batch_for_pack.remote(items_pack, pack_max_length))
+        
+        self.assertEqual(len(batches), 2)
+        self.assertEqual(len(batches[0]), 2) # 40 + 30 = 70
+        self.assertEqual(len(batches[1]), 1) # 40
+        self.assertEqual(batches[0][0]["seq_ctx"].input_ids.shape[1] + batches[0][1]["seq_ctx"].input_ids.shape[1], 70)
+        self.assertEqual(batches[1][0]["seq_ctx"].input_ids.shape[1], 40)
+        # --- Test 3: _pad_and_pack_batches ---
+        print("Testing _pad_and_pack_batches...")
+        # Input: First batch with length 70. Should pad 30 to reach 100. Second batch with length 40, should pad 60 to reach 100.
+        for idx, batch4pack_list in enumerate(batches):
+            packed_item = ray.get(train_controller._pad_and_pack_batches.remote(batch4pack_list, pack_max_length))
+            # Check total length
+            self.assertEqual(packed_item["seq_ctx"].input_ids.shape[1], pack_max_length)
+            # idx == 0:
+            if idx == 0:
+                # Check cu_seq_lens_q: [0, 40, 70, 100]
+                expected_cu_lens = torch.tensor([0, 40, 70, 100], dtype=torch.int32)
+                self.assertTrue(torch.equal(packed_item["seq_ctx"].cu_seq_lens_q, expected_cu_lens))
+                # Check padding labels are -100
+                self.assertTrue(torch.all(packed_item["shifted_labels"][0, 70:] == -100))
+            if idx == 1:
+                # Check cu_seq_lens_q: [0, 40, 100]
+                expected_cu_lens = torch.tensor([0, 40, 100], dtype=torch.int32)
+                self.assertTrue(torch.equal(packed_item["seq_ctx"].cu_seq_lens_q, expected_cu_lens))
+                # Check padding labels are -100
+                self.assertTrue(torch.all(packed_item["shifted_labels"][0, 40:] == -100))
+
+        # --- Test 4: _pad_to_max_packs_across_workes ---
+        pack_dummy = {"dummy": "pack"}
+        packed_data_batches = [
+            [[pack_dummy, pack_dummy]], # Worker 0: 2 packs
+            [[pack_dummy]]              # Worker 1: 1 pack
+        ]
+        # Execute the function locally
+        packed_data_batches = ray.get(train_controller._pad_to_max_packs_across_workes.remote(
+            packed_data_batches, 0, 2, pack_max_length
+        ))
+        # Verification
+        # Worker 0 should still have 2 packs
+        self.assertEqual(len(packed_data_batches[0][0]), 2)
+        
+        # Worker 1 should now have 2 packs (1 original + 1 padding)
+        self.assertEqual(len(packed_data_batches[1][0]), 2)
+        
+        # Verify the added item is a padding item
+        added_pack = packed_data_batches[1][0][1]
+        # Since we used the real _create_padding_item, it should have the correct structure
+        self.assertIn("seq_ctx", added_pack)
+        self.assertIn("shifted_labels", added_pack)
+        self.assertEqual(added_pack["seq_ctx"].input_ids.shape[1], pack_max_length)
+        self.assertTrue(torch.all(added_pack["shifted_labels"] == -100))
+        print("All controller logic tests passed!")
+
+if __name__ == "__main__":
+    unittest.main()

xtuner/v1/rl/base/__init__.py

@@ -1,6 +1,13 @@
 from .controller import TrainingController, TrainingControllerProxy, TrainingStepTimeLog
 from .loss import BaseRLLossConfig, RLLossContextInputItem
-from .worker import TrainingWorker, TrainingWorkerClass, TrainingWorkerProxy, WorkerConfig, WorkerLogItem
+from .worker import (
+    TrainingWorker,
+    TrainingWorkerClass,
+    TrainingWorkerProxy,
+    WorkerConfig,
+    WorkerInputItem,
+    WorkerLogItem,
+)
 
 
 __all__ = [
@@ -14,4 +21,5 @@
     "RLLossContextInputItem",
     "WorkerLogItem",
     "TrainingStepTimeLog",
+    "WorkerInputItem",
 ]

xtuner/v1/rl/base/controller.py

@@ -17,10 +17,10 @@
 from xtuner.v1.train.trainer import LoadCheckpointConfig
 from xtuner.v1.utils import get_logger, ray_method
 
+from .worker import TrainingWorker, WorkerInputItem, WorkerLogItem
 
-TRAIN_RAY_GET_TIMEOUT = os.getenv("XTUNER_TRAIN_RAY_GET_TIMEOUT", 5 * 3600)  # default 5 hours
 
-from .worker import TrainingWorker, WorkerInputItem, WorkerLogItem
+TRAIN_RAY_GET_TIMEOUT = os.getenv("XTUNER_TRAIN_RAY_GET_TIMEOUT", 5 * 3600)  # default 5 hours
 
 
 class TrainingStepTimeLog(TypedDict):
@@ -314,7 +314,10 @@ def _set_data_batches_properties(self, data_batches: list[WorkerInputItem]):
     def _pad_and_pack_batches(self, batch4pack: list[WorkerInputItem], pack_max_length: int) -> WorkerInputItem:
         seq_ctx_list = [item["seq_ctx"] for item in batch4pack]
         label_list = [item["shifted_labels"] for item in batch4pack]
-        advantage_list = [torch.tensor([item["advantages"]]).float().unsqueeze(0) for item in batch4pack]
+        advantage_list = []
+        for item in batch4pack:
+            advantages = item["advantages"].reshape(1, -1)
+            advantage_list.append(advantages)
         rollout_logprobs_list = [
             item["rollout_logprobs"] if self.has_rollout_logprobs else None for item in batch4pack
         ]
@@ -366,6 +369,7 @@ def _pad_to_max_packs_across_workes(
                 padding_item = self._create_padding_item(pack_max_length, pack_max_length)
                 padding_items = [padding_item for _ in range(num_padding_packs)]
                 packed_data_batches[dp_rank][step_idx].extend(padding_items)
+        return packed_data_batches
 
     @ray_method
     def fit(
@@ -428,7 +432,9 @@ def fit(
         # padding for each worker to have same number of packs in each optimizer step
         for step_idx in range(optimizer_steps):
             max_packs = max_packs_per_step[step_idx]
-            self._pad_to_max_packs_across_workes(packed_data_batches, step_idx, max_packs, pack_max_length)
+            packed_data_batches = self._pad_to_max_packs_across_workes(
+                packed_data_batches, step_idx, max_packs, pack_max_length
+            )
 
         pack_end_time = time.perf_counter()
         self.logger.info(f"Data packing took {pack_end_time - pack_start_time:.2f} seconds.")

xtuner/v1/rl/base/worker.py

@@ -516,13 +516,13 @@ def _apply_rollout_is_correction(
                 all_rollout_is_metrics.append(rollout_is_metrics)
                 all_mismatch_metrics.append(mismatch_metrics)
 
-        worker_log_item: WorkerLogItem = {"train_entropy": 0.0, "train_metrics": [], "sft_train_metrics": {}}
-        logger_msg = f"Rollout {rollout_idx}: "
-        sum_entropy = cast(torch.Tensor, sum_entropy)
-        dist.all_reduce(sum_entropy, op=dist.ReduceOp.SUM)
-        avg_sum_entropy = sum_entropy / global_grad_tokens if global_grad_tokens > 0 else torch.tensor(0.0)
-        worker_log_item["train_entropy"] = avg_sum_entropy.item()
-        logger_msg += f"avg entropy: {avg_sum_entropy:.4f}"
+        metrics = {
+            "sum_entropy": sum_entropy,
+            "sum_rollout_entropy": sum_rollout_entropy,
+            "all_mismatch_metrics": all_mismatch_metrics,
+            "all_rollout_is_metrics": all_rollout_is_metrics,
+        }
+        return loss_ctx_input_list, metrics
 
     @ray_method
     def fit(self, data_batches: list[list[WorkerInputItem]], rollout_idx: int):
@@ -579,10 +579,7 @@ def fit(self, data_batches: list[list[WorkerInputItem]], rollout_idx: int):
         global_grad_tokens = rank_grad_tokens.clone()
         dist.all_reduce(global_grad_tokens, op=dist.ReduceOp.SUM)
 
-        worker_log_item: WorkerLogItem = {
-            "train_entropy": 0.0,
-            "train_metrics": [],
-        }
+        worker_log_item: WorkerLogItem = {"train_entropy": 0.0, "train_metrics": [], "sft_train_metrics": {}}
         log_parts = []
 
         sum_entropy = cast(torch.Tensor, metrics["sum_entropy"])
@@ -678,7 +675,10 @@ def fit(self, data_batches: list[list[WorkerInputItem]], rollout_idx: int):
                 f"{key}={value:.4f}" if isinstance(value, float) else f"{key}={value}"
                 for key, value in log_info.items()
             )
-            log_str = f"Rank{self.rank} Rollout {rollout_idx} Step {i}: gradient_accumulation_steps={num_packs_this_step}" + log_str
+            log_str = (
+                f"Rank{self.rank} Rollout {rollout_idx} Step {i}: gradient_accumulation_steps={num_packs_this_step}, "
+                + log_str
+            )
             self.logger.info(log_str)
 
         self._rollout_step += 1

xtuner/v1/train/rl_trainer.py

@@ -33,6 +33,7 @@
     TrainingWorkerClass,
     TrainingWorkerProxy,
     WorkerConfig,
+    WorkerInputItem,
     WorkerLogItem,
 )
 from xtuner.v1.rl.base import TrainingWorker as BaseTrainingWorker
@@ -769,10 +770,10 @@ def _prepare_train_data(self, data_groups, pack_max_length, multimodal_train_inf
                     rollout_logprobs = None
 
                 seq_ctx = get_train_seq_ctx(input_ids, multimodal_train_info, len(response_ids) - 1)
-                data_dict = {
+                data_dict: WorkerInputItem = {
                     "seq_ctx": seq_ctx,
                     "shifted_labels": shifted_labels,
-                    "advantage": advantages[i].item(),
+                    "advantages": advantages[i],
                     "rollout_logprobs": rollout_logprobs,
                 }