reduce consumed samples out of data_time (#1389)

Author: jayhenry

xtuner/v1/ray/judger/native.py

@@ -3,7 +3,7 @@
 
 import httpx
 import ray
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 from ray.util.placement_group import PlacementGroup
 
 from xtuner.v1.data_proto.rl_data import RLDataFlowItem, RLJudgerResponseItem
@@ -13,7 +13,7 @@
 class NativeJudgerConfig(BaseModel):
     """Base configuration class for judgers."""
 
-    model_config = {"arbitrary_types_allowed": True}
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
     judger_name: str
     num_ray_actors: int = 1
     reward_func: Optional[Callable] = None

xtuner/v1/train/trainer.py

@@ -706,6 +706,7 @@ def fit(self):
 
             time_before_train_step = time.time()
             data_time = time_before_train_step - time_before_get_data
+            cur_sample_num = len(data_batch)
 
             seq_ctx_list: list[SequenceContext] = []
             loss_ctx_input_list: list[CELossContextInputItem] = []
@@ -778,6 +779,7 @@ def fit(self):
 
             self._cur_step += 1
 
+            self._reduced_consumed_samples += self._reduce_number_across_rank(cur_sample_num)
             reduced_step_consumed_tokens = self._reduce_number_across_rank(step_consumed_tokens)
             self._reduced_consumed_tokens += reduced_step_consumed_tokens
 
@@ -1257,7 +1259,6 @@ def _data_iter(self):
                 data_iter = iter(self._dataloader)
                 data = next(data_iter)
 
-            self._reduced_consumed_samples += self._reduce_number_across_rank(len(data))
             yield data
 
     def _get_checkpoint_path(self, epoch: int, step: int, is_snapshot: bool = False) -> Path: