_save_dataloader return total_consumed_samples

Author: jayhenry

xtuner/v1/datasets/dataloader.py

@@ -57,12 +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 get_total_consumed_samples(self) -> int:
-        sampler = self.sampler
-        if hasattr(sampler, "get_total_consumed_steps"):
-            return int(sampler.get_total_consumed_steps())
-        return 0
-
     # __iter__ is inherited from torch.utils.data.DataLoader
 
     # Streaming dataloader may not have `set_epoch` and `__len__` method, so we add here.

xtuner/v1/datasets/preset_sampler.py

@@ -174,9 +174,6 @@ def __len__(self) -> int:
     def set_epoch(self, epoch: int) -> None:
         self.epoch = epoch
 
-    def get_total_consumed_steps(self) -> int:
-        return self._consumed.total_for_checkpoint()
-
     def get_state_dict(self, step: int | None = None) -> dict:
         # Same convention as :class:`LengthGroupedSampler`: ``step`` is the global pack offset
         # (modulo ``total_size``) into ``global_order``, shared across all ranks in the checkpoint.

xtuner/v1/datasets/sampler.py

@@ -135,9 +135,6 @@ def set_epoch(self, epoch: int) -> None:
         """
         self.epoch = epoch
 
-    def get_total_consumed_steps(self) -> int:
-        return self._consumed.total_for_checkpoint()
-
     def load_state_dict(self, state_dict) -> None:
         """Load the sampler state.
 
@@ -295,9 +292,6 @@ def set_epoch(self, epoch: int) -> None:
         """
         self.epoch = epoch
 
-    def get_total_consumed_steps(self) -> int:
-        return self._consumed.total_for_checkpoint()
-
     def load_state_dict(self, state_dict: dict) -> None:
         """Load the sampler state.
 

xtuner/v1/train/trainer.py

@@ -1129,10 +1129,9 @@ def _maybe_save(self, is_snapshot: bool = False) -> bool:
         total_consumed_tokens = (
             self._reduce_number_across_rank(self._local_total_consumed_tokens) + self._init_total_tokens
         )
-        total_consumed_samples = self._dataloader.get_total_consumed_samples()
 
         # Save dataloader
-        self._save_dataloader(dataloader_path)
+        total_consumed_samples = self._save_dataloader(dataloader_path)
 
         DEVICE_MODULE.empty_cache()
 
@@ -1211,10 +1210,11 @@ def _maybe_save(self, is_snapshot: bool = False) -> bool:
 
         return True
 
-    def _save_dataloader(self, dataloader_path: Path | str):
+    def _save_dataloader(self, dataloader_path: Path | str) -> int:
         dataloader_state = self._dataloader.get_state_dict()
         if self.rank == 0:
             torch.save(dataloader_state, dataloader_path)
+        return int(dataloader_state.get("sampler", {}).get("total_consumed_steps", 0))
 
     @property
     def work_dir(self) -> Path: