add ut and unify training input to WorkerLogItem

Author: YanhuiDua

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
@@ -774,10 +775,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,
                 }