refine code

Author: jayhenry

xtuner/v1/datasets/dataloader.py

@@ -54,6 +54,7 @@ def get_state_dict(self, consumed_samples: int = -1) -> dict:
                 "Dataloader.get_state_dict(consumed_samples=...) is deprecated; use the default (-1). "
                 "Consumed samples are tracked on the sampler."
             )
+        # TODO: remove consumed_samples parameter in get_dataloader_state in next major release
         dataloader_state = get_dataloader_state(self, consumed_samples)
         return cast(dict, dataloader_state)
 

xtuner/v1/datasets/preset_sampler.py

@@ -192,11 +192,8 @@ def get_state_dict(self, step: int | None = None) -> dict:
         }
 
     def load_state_dict(self, state_dict: dict) -> None:
-        tc = state_dict.get("total_consumed_steps")
-        if tc is not None:
-            self._consumed.set_init_from_checkpoint(int(tc))
-        else:
-            self._consumed.set_init_from_checkpoint(0)
+        tc = int(state_dict.get("total_consumed_steps", 0))
+        self._consumed.set_init_from_checkpoint(tc)
         if self.world_size != state_dict.get("world_size"):
             logger.warning(
                 f"PresetSampler: world_size mismatch: checkpoint has "

xtuner/v1/datasets/resume.py

@@ -5,6 +5,7 @@
 
 from .consumed_steps import apply_old_ckpt_init_steps
 from .packing import ExpandSoftPackDataset, _LegacySoftPackDataset
+from .preset_sampler import PresetSampler
 from .sampler import LengthGroupedSampler, ParallelSampler
 
 
@@ -45,7 +46,7 @@ def load_dataloader_state(
     state: dict,
     train_state_total_consumed_samples: int | None = None,
 ):
-    sampler = dataloader.sampler
+    sampler: ParallelSampler | LengthGroupedSampler | PresetSampler = dataloader.sampler  # type: ignore[assignment]
     dataset = dataloader.dataset
 
     # Sampler require `load_state_dict` to restore the training progress since the sampler state will

xtuner/v1/datasets/sampler.py

@@ -141,11 +141,8 @@ def load_state_dict(self, state_dict) -> None:
         Args:
             state_dict (dict): The state of the sampler.
         """
-        tc = state_dict.get("total_consumed_steps")
-        if tc is not None:
-            self._consumed.set_init_from_checkpoint(int(tc))
-        else:
-            self._consumed.set_init_from_checkpoint(0)
+        tc = int(state_dict.get("total_consumed_steps", 0))
+        self._consumed.set_init_from_checkpoint(tc)
         self.epoch = state_dict["epoch"]
         self.step = state_dict["step"]
 
@@ -298,11 +295,8 @@ def load_state_dict(self, state_dict: dict) -> None:
         Args:
             state_dict (dict): The state of the sampler.
         """
-        tc = state_dict.get("total_consumed_steps")
-        if tc is not None:
-            self._consumed.set_init_from_checkpoint(int(tc))
-        else:
-            self._consumed.set_init_from_checkpoint(0)
+        tc = int(state_dict.get("total_consumed_steps", 0))
+        self._consumed.set_init_from_checkpoint(tc)
         self.epoch = state_dict["epoch"]
         self.step = state_dict["step"]