[Enhance] enhance logging (mostly on token stats)

1. use reduced tokens stats for exp_tgs;
2. rename some variables (drop the reduced_ prefix, use total_ instead
for clarity);
3. drop maxvio stats in loss_log as it is already covered in internal
metrics.

Author: nil0x9

xtuner/v1/engine/train_engine.py

@@ -53,7 +53,7 @@ class LossLog(TypedDict):
 class OtherLog(TypedDict):
     __pydantic_config__ = ConfigDict(arbitrary_types_allowed=True)  # type: ignore[misc]
     maxvio: NotRequired[float]
-    consumed_tokens: float
+    consumed_tokens: int
     extra_info: ModelForwardExtraLogInfo
     efficient_attn_ratio: float
 
@@ -252,7 +252,7 @@ def train_step(self, data_batches: list[ModelItem]) -> tuple[LossLog, OtherLog]:
         step_llm_loss = torch.tensor(0.0, device=DEVICE)
         step_balancing_loss: torch.Tensor | None = None
         step_z_loss: torch.Tensor | None = None
-        step_consumed_tokens = torch.tensor(0.0, device=DEVICE)
+        step_consumed_tokens = torch.tensor(0, device=DEVICE)
 
         if self._count == 0:
             logger.info(f"grad_accumulation_steps: {iters_per_step}")
@@ -346,7 +346,7 @@ def train_step(self, data_batches: list[ModelItem]) -> tuple[LossLog, OtherLog]:
             reduced_z_loss = step_z_loss
             dist.all_reduce(reduced_z_loss.div_(dist.get_world_size()))
             loss_log["reduced_z_loss"] = reduced_z_loss.item()
-        other_log["consumed_tokens"] = step_consumed_tokens.item()
+        other_log["consumed_tokens"] = cast(int, step_consumed_tokens.item())
         other_log["extra_info"] = train_engine_extra_info
         other_log["efficient_attn_ratio"] = (efficient_forward_tokens / total_forward_tokens).item()
         return loss_log, other_log

xtuner/v1/engine/vision_compose_train_engine.py

@@ -183,7 +183,7 @@ def train_step(self, data_batches: List[ModelItem]) -> tuple[LossLog, OtherLog]:
         step_llm_loss = torch.tensor(0.0, device=DEVICE)
         step_balancing_loss: torch.Tensor | None = None
         step_z_loss: torch.Tensor | None = None
-        step_consumed_tokens = torch.tensor(0.0, device=DEVICE)
+        step_consumed_tokens = torch.tensor(0, device=DEVICE)
         efficient_forward_tokens = torch.tensor(0, device=DEVICE, dtype=torch.long)
         total_forward_tokens = torch.tensor(0, device=DEVICE, dtype=torch.long)
 
@@ -257,7 +257,7 @@ def train_step(self, data_batches: List[ModelItem]) -> tuple[LossLog, OtherLog]:
             reduced_z_loss = step_z_loss
             dist.all_reduce(reduced_z_loss.div_(dist.get_world_size()))
             loss_log["reduced_z_loss"] = reduced_z_loss.item()
-        other_log["consumed_tokens"] = step_consumed_tokens.item()
+        other_log["consumed_tokens"] = cast(int, step_consumed_tokens.item())
         other_log["extra_info"] = train_engine_extra_info  # type: ignore[assignment]
         other_log["efficient_attn_ratio"] = (efficient_forward_tokens / total_forward_tokens).item()
         return loss_log, other_log

xtuner/v1/train/trainer.py

@@ -514,9 +514,9 @@ def __init__(
         self._debug = debug
         self._seed = seed
 
-        self._reduced_consumed_tokens = 0
+        self._total_consumed_tokens = 0
         self._exp_consumed_tokens = 0
-        self._reduced_consumed_samples = 0
+        self._total_consumed_samples = 0
 
         self._train_time = 0
         self._train_time_offset = 0
@@ -762,26 +762,24 @@ def fit(self):
                 extra_info_dict = extra_info_updated.get()
             loss_log.update(extra_info_dict)
 
-            if "maxvio" in other_log:
-                loss_log["maxvio"] = other_log["maxvio"]
             loss_log["efficient_attn_ratio"] = other_log["efficient_attn_ratio"]
 
             internal_metrics = self._maybe_pop_model_internal_metrics(engine_input)
 
             self._cur_step += 1
 
             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._total_consumed_tokens += reduced_step_consumed_tokens
+            self._exp_consumed_tokens += reduced_step_consumed_tokens
             self._train_time = time_after_train_step - train_begin
 
             # TODO: This log should be move before lr_scheduler.step, but for CI BC, keep it temporarily
             self._log_step(
                 loss_log=loss_log,
-                step_consumed_tokens=step_consumed_tokens,
+                local_step_consumed_tokens=step_consumed_tokens,
+                step_consumed_tokens=reduced_step_consumed_tokens,
                 exp_consumed_tokens=self._exp_consumed_tokens,
-                reduced_consumed_tokens=self._reduced_consumed_tokens,
+                total_consumed_tokens=self._total_consumed_tokens,
                 data_time=data_time,
                 step_time=step_time,
                 train_time=self._train_time,
@@ -1137,8 +1135,8 @@ def _maybe_save(self, is_snapshot: bool = False) -> bool:
                         {
                             "cur_step": self.cur_step,
                             "cur_epoch": self._cur_epoch,
-                            "reduced_consumed_samples": self._reduced_consumed_samples,
-                            "reduced_consumed_tokens": self._reduced_consumed_tokens,
+                            "total_consumed_samples": self._total_consumed_samples,
+                            "total_consumed_tokens": self._total_consumed_tokens,
                             "train_time_offset": self._train_time + self._train_time_offset,
                         }
                     )
@@ -1150,8 +1148,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 = int(self._reduced_consumed_samples)
-        current_exp.consumed_tokens = int(self._reduced_consumed_tokens)
+        current_exp.consumed_samples = int(self._total_consumed_samples)
+        current_exp.consumed_tokens = int(self._total_consumed_tokens)
         current_exp.history[-1]["end"] = self.cur_step
 
         # Delete checkpoints and update meta's checkpoint_list
@@ -1188,7 +1186,7 @@ def _maybe_save(self, is_snapshot: bool = False) -> bool:
 
     def _save_dataloader(self, dataloader_path: Path | str):
         if self.rank == 0:
-            dataloader_state = self._dataloader.get_state_dict(self._reduced_consumed_samples)
+            dataloader_state = self._dataloader.get_state_dict(self._total_consumed_samples)
             torch.save(dataloader_state, dataloader_path)
 
     @property
@@ -1243,7 +1241,7 @@ def _data_iter(self):
                 data_iter = iter(self._dataloader)
                 data = next(data_iter)
 
-            self._reduced_consumed_samples += self._reduce_number_across_rank(len(data))
+            self._total_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:
@@ -1413,10 +1411,11 @@ def _maybe_profiling(self):
 
     def _log_step(
         self,
-        loss_log: dict,
+        loss_log: LossLog,
+        local_step_consumed_tokens: int,
         step_consumed_tokens: int,
         exp_consumed_tokens: int,
-        reduced_consumed_tokens: int,
+        total_consumed_tokens: int,
         data_time: float,
         step_time: float,
         train_time: float,
@@ -1426,20 +1425,20 @@ def _log_step(
     ):
         """Log the training step information."""
         e2e_train_time = train_time + train_time_offset
-        tgs = step_consumed_tokens / step_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
+        total_consumed_tokens_per_rank = total_consumed_tokens / self.world_size
+        exp_consumed_tokens_per_rank = exp_consumed_tokens / self.world_size
+
+        tgs = local_step_consumed_tokens / step_time
+        e2e_tgs = total_consumed_tokens_per_rank / e2e_train_time
+        exp_tgs = exp_consumed_tokens_per_rank / train_time
         lr = self._lr_scheduler.get_last_lr()[0]
 
         remaining_steps = self.total_step - self.cur_step
-        avg_tokens_per_step = rank_consumed_tokens / self.cur_step
+        avg_tokens_per_step = total_consumed_tokens_per_rank / 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)))
 
-        est_global_batch_tokens = self.data_mesh["dp"].size() * step_consumed_tokens
-
         loss_log_list = [f"{k}: {v:.8f}" for k, v in loss_log.items()]
         loss_log_str = ", ".join(loss_log_list)
 
@@ -1453,16 +1452,16 @@ def _log_step(
         self.logger.info(
             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"reduced_consumed_tokens: {reduced_consumed_tokens} "
+            f"text_tokens: {local_step_consumed_tokens} "
+            f"step_consumed_tokens: {step_consumed_tokens} "
+            f"total_consumed_tokens: {total_consumed_tokens} "
             f"{loss_log_str} "
             f"grad_norm: {grad_norm:.8f} "
             f"max_memory: {max_memory / (1024**3):.2f} GB "
             f"reserved_memory: {reserved_memory / (1024**3):.2f} GB "
             f"tgs: {tgs:.1f} "
             f"exp_tgs: {exp_tgs: .1f} "
             f"e2e_tgs: {e2e_tgs:.1f} "
-            f"est_global_batch_tokens: {est_global_batch_tokens} "
             f"eta: {eta_hms} "
         )
 
@@ -1472,9 +1471,9 @@ def _log_step(
             "time/step_time": round(step_time, 4),
             "time/train_time": round(train_time, 4),
             "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/reduced_consumed_tokens": reduced_consumed_tokens,
+            "runtime_info/text_tokens": local_step_consumed_tokens,
+            "runtime_info/step_consumed_tokens": step_consumed_tokens,
+            "runtime_info/total_consumed_tokens": total_consumed_tokens,
             "runtime_info/tgs": tgs,
             "runtime_info/exp_tgs": exp_tgs,
             "runtime_info/e2e_tgs": e2e_tgs,
@@ -1679,13 +1678,13 @@ def _load_checkpoint(self):
         self._cur_epoch = train_state["cur_epoch"]
 
         if load_checkpoint_cfg.load_dataset:
-            self._reduced_consumed_tokens = train_state.get("reduced_consumed_tokens", 0)  # default 0 for BC
+            self._total_consumed_tokens = train_state.get("total_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
+            # _total_consumed_samples 会影响 save dcp时 dataloader.get_state_dict的状态。
+            # 1) 如果加载 dataset，应该恢复_total_consumed_samples为checkpoint中的值。
+            # 2) 如果不加载 dataset，应该保持_total_consumed_samples为初始值0，否则如果加载上旧dataloader的total_consumed_samples
+            #    会导致存储新dataloader时 total_consumed_samples 是不正确的值。
+            self._total_consumed_samples = train_state.get("total_consumed_samples", 0)  # default 0 for BC
 
             dataloader_path = resume_from / self._SAVE_DATALOADER_DIR
             self._resume_dataloader(dataloader_path)