refine local_steps update logic

Author: jayhenry

xtuner/v1/datasets/dataloader.py

@@ -57,10 +57,6 @@ def get_state_dict(self, consumed_samples: int = -1) -> dict:
         dataloader_state = get_dataloader_state(self, consumed_samples)
         return cast(dict, dataloader_state)
 
-    def record_consumed_samples(self, n: int) -> None:
-        if hasattr(self.sampler, "record_consumed_samples"):
-            self.sampler.record_consumed_samples(n)
-
     def get_total_consumed_samples(self) -> int:
         sampler = self.sampler
         if hasattr(sampler, "get_total_consumed_steps"):

xtuner/v1/datasets/preset_sampler.py

@@ -163,7 +163,9 @@ def __init__(
     def __iter__(self) -> Iterator[int]:
         # load order from npy → global_order → rank_view 类型均为 memmap, 子视图 的路径仍然保持
         # memmap 语义（视图、按需分页、文件后端）；单机多进程可共享同一份文件页缓存
-        yield from self.global_order[self.step + self.rank : self.total_size : self.world_size]
+        for idx in self.global_order[self.step + self.rank : self.total_size : self.world_size]:
+            self._consumed.record(1)
+            yield int(idx)
         self.step = 0
 
     def __len__(self) -> int:
@@ -172,9 +174,6 @@ def __len__(self) -> int:
     def set_epoch(self, epoch: int) -> None:
         self.epoch = epoch
 
-    def record_consumed_samples(self, n: int) -> None:
-        self._consumed.record(n)
-
     def get_total_consumed_steps(self) -> int:
         return self._consumed.total_for_checkpoint()
 

xtuner/v1/datasets/sampler.py

@@ -113,7 +113,9 @@ def __iter__(self) -> Iterator[int]:
         # subsample
         indices = indices[self.step + self.rank : self.total_size : self.world_size]
 
-        yield from iter(indices)
+        for idx in indices:
+            self._consumed.record(1)
+            yield idx
         self.step = 0
 
     def __len__(self) -> int:
@@ -133,9 +135,6 @@ def set_epoch(self, epoch: int) -> None:
         """
         self.epoch = epoch
 
-    def record_consumed_samples(self, n: int) -> None:
-        self._consumed.record(n)
-
     def get_total_consumed_steps(self) -> int:
         return self._consumed.total_for_checkpoint()
 
@@ -275,7 +274,9 @@ def __iter__(self) -> Iterator[int]:
         assert len(indices) == self.total_size
         indices = indices[self.step + self.rank : self.total_size : self.world_size]
         assert len(indices) == self.num_samples - self.step // self.world_size
-        yield from iter(indices)
+        for idx in indices:
+            self._consumed.record(1)
+            yield idx
         self.step = 0
 
     def __len__(self) -> int:
@@ -294,9 +295,6 @@ def set_epoch(self, epoch: int) -> None:
         """
         self.epoch = epoch
 
-    def record_consumed_samples(self, n: int) -> None:
-        self._consumed.record(n)
-
     def get_total_consumed_steps(self) -> int:
         return self._consumed.total_for_checkpoint()
 

xtuner/v1/train/trainer.py

@@ -727,7 +727,6 @@ def fit(self):
         train_begin = time.time()
         time_before_get_data = time.time()
         for data_batch in self._data_iter():
-            consumed_samples = len(data_batch)
             time_before_train_step = time.time()
 
             ProberList.set_step(self._cur_step + 1)
@@ -762,7 +761,6 @@ def fit(self):
             self._cur_step += 1
             step_tokens = train_step_info["step_consumed_tokens"]
             self._local_total_consumed_tokens += step_tokens
-            self._dataloader.record_consumed_samples(consumed_samples)
             self._train_time = time_after_train_step - train_begin
 
             # Compute training metrics