add trainingStepTimeLog in train controller

Author: YanhuiDua

xtuner/v1/rl/base/__init__.py

@@ -1,4 +1,4 @@
-from .controller import TrainingController, TrainingControllerProxy
+from .controller import TrainingController, TrainingControllerProxy, TrainingStepTimeLog
 from .loss import BaseRLLossConfig, RLLossContextInputItem
 from .worker import TrainingWorker, TrainingWorkerClass, TrainingWorkerProxy, WorkerConfig, WorkerLogItem
 
@@ -13,4 +13,5 @@
     "BaseRLLossConfig",
     "RLLossContextInputItem",
     "WorkerLogItem",
+    "TrainingStepTimeLog",
 ]

xtuner/v1/rl/base/controller.py

@@ -1,13 +1,15 @@
 import math
 import os
 import random
+import time
 from pathlib import Path
 from typing import Literal, cast
 
 import numpy as np
 import ray
 import torch
 from ray.actor import ActorProxy
+from typing_extensions import TypedDict
 
 from xtuner.v1.data_proto.sequence_context import SequenceContext
 from xtuner.v1.model.compose.base import BaseComposeConfig
@@ -21,6 +23,11 @@
 from .worker import TrainingWorker, WorkerInputItem, WorkerLogItem
 
 
+class TrainingStepTimeLog(TypedDict):
+    data_packing_time: float
+    worker_training_time: float
+
+
 class RawTrainingController:
     def __init__(self, workers: list[TrainingWorker]) -> None:
         self.workers = workers
@@ -367,7 +374,8 @@ def fit(
         pack_max_length: int,
         rollout_idx: int,
         enable_dp_balance: bool = True,
-    ) -> list[WorkerLogItem]:
+    ) -> tuple[list[WorkerLogItem], TrainingStepTimeLog]:
+        pack_start_time = time.perf_counter()
         self._set_data_batches_properties(data_batches)
 
         world_size = len(self.workers)
@@ -422,6 +430,9 @@ def fit(
             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)
 
+        pack_end_time = time.perf_counter()
+        self.logger.info(f"Data packing took {pack_end_time - pack_start_time:.2f} seconds.")
+
         handles = []
         for worker_idx, worker in enumerate(self.workers):
             handles.append(
@@ -430,8 +441,14 @@ def fit(
                     rollout_idx=rollout_idx,
                 )
             )
-        log_infos = ray.get(handles, timeout=TRAIN_RAY_GET_TIMEOUT)
-        return log_infos
+        worker_log_infos = ray.get(handles)
+        fit_end_time = time.perf_counter()
+        self.logger.info(f"Training step took {fit_end_time - pack_end_time:.2f} seconds.")
+        training_time: TrainingStepTimeLog = {
+            "data_packing_time": pack_end_time - pack_start_time,
+            "worker_training_time": fit_end_time - pack_end_time,
+        }
+        return worker_log_infos, training_time
 
     @ray_method
     def offload(self, target: Literal["model", "optimizer", "all"] = "all"):

xtuner/v1/train/rl_trainer.py

@@ -29,6 +29,7 @@
 from xtuner.v1.rl.base import (
     TrainingController,
     TrainingControllerProxy,
+    TrainingStepTimeLog,
     TrainingWorkerClass,
     TrainingWorkerProxy,
     WorkerConfig,
@@ -555,12 +556,18 @@ def _train_step(self, rollout_idx: int, data_groups, multimodal_train_infos, ste
         )
 
         with timer("training", step_timer_dict):
-            workers_log_item: List[WorkerLogItem] = ray.get(
+            workers_log_item: List[WorkerLogItem]
+            training_time: TrainingStepTimeLog
+            workers_log_item, training_time = ray.get(
                 self._train_controller.fit.remote(
                     data_batches, pack_max_length=self._train_worker_cfg.pack_max_length, rollout_idx=rollout_idx
                 )
             )
         self._writer.add_scalar(tag="time/training", scalar_value=step_timer_dict["training"], global_step=rollout_idx)
+        self._writer.add_scalars(
+            tag_scalar_dict={f"time/train_{key}": cast(float, value) for key, value in training_time.items()},
+            global_step=rollout_idx,
+        )
 
         rank0_log_item = workers_log_item[0]
         # These metrics are already aggregated across distributed workers and logging only the metrics from rank 0.