1) Add reduced consumed tokens/samples. 2) Do not resume consumed samples if not load dataset (#1326)

* [Feature] add reduced consumed tokens and samples
* [Fix] Do not resume consumed samples and tokens if not load_dataset

Author: jayhenry

xtuner/v1/train/trainer.py

@@ -514,9 +514,9 @@ def __init__(
         self._debug = debug
         self._seed = seed
 
-        self._consumed_tokens = 0
+        self._reduced_consumed_tokens = 0
         self._exp_consumed_tokens = 0
-        self._consumed_samples = 0
+        self._reduced_consumed_samples = 0
 
         self._train_time = 0
         self._train_time_offset = 0
@@ -769,7 +769,10 @@ def fit(self):
             internal_metrics = self._maybe_pop_model_internal_metrics(engine_input)
 
             self._cur_step += 1
-            self._consumed_tokens += step_consumed_tokens
+
+            reduced_step_consumed_tokens = self._reduce_number_across_rank(step_consumed_tokens)
+            self._reduced_consumed_tokens += reduced_step_consumed_tokens
+
             self._exp_consumed_tokens += step_consumed_tokens
             self._train_time = time_after_train_step - train_begin
 
@@ -778,7 +781,7 @@ def fit(self):
                 loss_log=loss_log,
                 step_consumed_tokens=step_consumed_tokens,
                 exp_consumed_tokens=self._exp_consumed_tokens,
-                total_consumed_tokens=self._consumed_tokens,
+                reduced_consumed_tokens=self._reduced_consumed_tokens,
                 data_time=data_time,
                 step_time=step_time,
                 train_time=self._train_time,
@@ -805,6 +808,12 @@ def fit(self):
             self._metrics_recorder.close()
         self.logger.info(f"Training finished in {time.time() - train_begin:.2f} seconds")
 
+    def _reduce_number_across_rank(self, rank_number: int) -> int:
+        _gathered_list = [None for _ in range(self.world_size)]
+        dist.all_gather_object(_gathered_list, rank_number)
+        reduced_number = sum(_gathered_list)  # type: ignore[arg-type]
+        return reduced_number
+
     def _maybe_init_model_metrics_recorder(
         self,
         internal_metrics_cfg: InternalMetricsConfig | None,
@@ -1128,8 +1137,8 @@ def _maybe_save(self, is_snapshot: bool = False) -> bool:
                         {
                             "cur_step": self.cur_step,
                             "cur_epoch": self._cur_epoch,
-                            "consumed_samples": self._consumed_samples,
-                            "consumed_tokens": self._consumed_tokens,
+                            "reduced_consumed_samples": self._reduced_consumed_samples,
+                            "reduced_consumed_tokens": self._reduced_consumed_tokens,
                             "train_time_offset": self._train_time + self._train_time_offset,
                         }
                     )
@@ -1141,8 +1150,8 @@ def _maybe_save(self, is_snapshot: bool = False) -> bool:
         ckp_list.append(str(checkpoint_path))
         current_exp.cur_step = self.cur_step
         current_exp.cur_epoch = self._cur_epoch
-        current_exp.consumed_samples = self._consumed_samples
-        current_exp.consumed_tokens = int(self._consumed_tokens)
+        current_exp.consumed_samples = int(self._reduced_consumed_samples)
+        current_exp.consumed_tokens = int(self._reduced_consumed_tokens)
         current_exp.history[-1]["end"] = self.cur_step
 
         # Delete checkpoints and update meta's checkpoint_list
@@ -1178,12 +1187,8 @@ def _maybe_save(self, is_snapshot: bool = False) -> bool:
         return True
 
     def _save_dataloader(self, dataloader_path: Path | str):
-        _gathered_list = [None for _ in range(self.data_mesh["dp"].size())]
-        dist.all_gather_object(_gathered_list, self._consumed_samples, group=self.data_mesh["dp"].get_group())
-        global_consumed_samples = sum(_gathered_list)  # type: ignore[arg-type]
-
         if self.rank == 0:
-            dataloader_state = self._dataloader.get_state_dict(global_consumed_samples)
+            dataloader_state = self._dataloader.get_state_dict(self._reduced_consumed_samples)
             torch.save(dataloader_state, dataloader_path)
 
     @property
@@ -1232,12 +1237,13 @@ def _data_iter(self):
             # dist.breakpoint(skip=14)
             try:
                 data = next(data_iter)
-                self._consumed_samples += len(data)
             except StopIteration:
                 self._cur_epoch += 1
                 self._dataloader.set_epoch(self._cur_epoch)
                 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:
@@ -1410,7 +1416,7 @@ def _log_step(
         loss_log: dict,
         step_consumed_tokens: int,
         exp_consumed_tokens: int,
-        total_consumed_tokens: int,
+        reduced_consumed_tokens: int,
         data_time: float,
         step_time: float,
         train_time: float,
@@ -1421,12 +1427,13 @@ def _log_step(
         """Log the training step information."""
         e2e_train_time = train_time + train_time_offset
         tgs = step_consumed_tokens / step_time
-        e2e_tgs = total_consumed_tokens / e2e_train_time
+        rank_consumed_tokens = reduced_consumed_tokens / self.world_size
+        e2e_tgs = rank_consumed_tokens / e2e_train_time
         exp_tgs = exp_consumed_tokens / train_time
         lr = self._lr_scheduler.get_last_lr()[0]
 
         remaining_steps = self.total_step - self.cur_step
-        avg_tokens_per_step = total_consumed_tokens / self.cur_step
+        avg_tokens_per_step = rank_consumed_tokens / self.cur_step
         remaining_tokens = remaining_steps * avg_tokens_per_step
         eta_seconds = remaining_tokens / (tgs + 1e-12)
         eta_hms = str(timedelta(seconds=int(eta_seconds)))
@@ -1447,7 +1454,7 @@ def _log_step(
             f"Epoch {self._cur_epoch} Step {self.cur_step}/{self.total_step} "
             f"data_time: {data_time:.4f} lr: {lr:.6e} time: {step_time:.4f} "
             f"text_tokens: {step_consumed_tokens} "
-            f"total_consumed_tokens: {total_consumed_tokens} "
+            f"reduced_consumed_tokens: {reduced_consumed_tokens} "
             f"{loss_log_str} "
             f"grad_norm: {grad_norm:.8f} "
             f"max_memory: {max_memory / (1024**3):.2f} GB "
@@ -1467,7 +1474,7 @@ def _log_step(
             "time/eta_seconds": round(eta_seconds, 1),
             "runtime_info/text_tokens": step_consumed_tokens,
             "runtime_info/est_global_batch_tokens": est_global_batch_tokens,
-            "runtime_info/total_consumed_tokens": total_consumed_tokens,
+            "runtime_info/reduced_consumed_tokens": reduced_consumed_tokens,
             "runtime_info/tgs": tgs,
             "runtime_info/exp_tgs": exp_tgs,
             "runtime_info/e2e_tgs": e2e_tgs,
@@ -1663,20 +1670,25 @@ def _load_checkpoint(self):
             load_args=load_checkpoint_cfg.load_optimizer_args,
         )
 
-        if load_checkpoint_cfg.load_dataset:
-            dataloader_path = resume_from / self._SAVE_DATALOADER_DIR
-            self._resume_dataloader(dataloader_path)
-
         train_state_path = resume_from / self._SAVE_TRAIN_STATE_PATH
 
         with train_state_path.open("r") as f:
             train_state = json.load(f)
 
         self._cur_step = train_state["cur_step"]
         self._cur_epoch = train_state["cur_epoch"]
-        self._consumed_samples = train_state["consumed_samples"]
-        self._consumed_tokens = train_state["consumed_tokens"]
-        self._train_time_offset = train_state["train_time_offset"]
+
+        if load_checkpoint_cfg.load_dataset:
+            self._reduced_consumed_tokens = train_state.get("reduced_consumed_tokens", 0)  # default 0 for BC
+            self._train_time_offset = train_state["train_time_offset"]
+            # _reduced_consumed_samples 会影响 save dcp时 dataloader.get_state_dict的状态。
+            # 1) 如果加载 dataset，应该恢复_reduced_consumed_samples为checkpoint中的值。
+            # 2) 如果不加载 dataset，应该保持_reduced_consumed_samples为初始值0，否则如果加载上旧dataloader的reduced_consumed_samples
+            #    会导致存储新dataloader时 reduced_consumed_samples 是不正确的值。
+            self._reduced_consumed_samples = train_state.get("reduced_consumed_samples", 0)  # default 0 for BC
+
+            dataloader_path = resume_from / self._SAVE_DATALOADER_DIR
+            self._resume_dataloader(dataloader_path)
 
         if load_checkpoint_cfg.load_scheduler:
             scheduler_path = resume_from / self._SAVE_SCHEDULER_DIR