[Enhance] Make ModelOutputs pydantic BaseModel

ghstack-source-id: 2ad7bfb
Pull-Request: InternLM#1516

Author: HAOCHENYE

xtuner/v1/engine/train_engine.py

@@ -293,6 +293,7 @@ def train_step(self, data_batches: list[ModelItem]) -> tuple[LossLog, OtherLog]:
                     seq_ctx=seq_ctx_list,
                     loss_ctx=loss_ctx_list,
                 )
+            output.free_nongrad_feature()
 
             # llm loss has been global averaged
             llm_loss = output["loss"]

xtuner/v1/loss/base_loss_ctx.py

@@ -186,6 +186,8 @@ def forward(
         head_weight: torch.Tensor,
         head_bias: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, tuple[torch.Tensor | None, dict[str, Any]]]:
+        from xtuner.v1.model.utils.misc import ModelForwardExtraLogInfo
+
         assert self.loss_kwargs is not None, "loss_kwargs must be set before calling forward"
         if head_bias is not None:
             raise NotImplementedError("Loss does not support head_bias yet.")
@@ -195,6 +197,10 @@ def forward(
         else:
             loss, (logits, extra_info) = self.chunk_mode(hidden_states, head_weight, head_bias, self.loss_kwargs)
 
+        # TODO: yanhuida, should be removed
+        if not isinstance(extra_info, ModelForwardExtraLogInfo):
+            extra_info = ModelForwardExtraLogInfo(extra_info)
+
         extra_info["local_base_loss"] = loss.detach().clone()
 
         # Step 2.c in the loss calculation: reduce the loss over all ranks using all_reduce with autograd support

xtuner/v1/model/base.py

@@ -194,11 +194,32 @@ def layers_type(self) -> list[Literal["full_attention", "sliding_attention"]]:
             ]
 
 
-class ModelOutputs(TypedDict):
-    hidden_states: NotRequired[list[torch.Tensor]]
-    logits: NotRequired[torch.Tensor]
+class ModelOutputs(PydanticBaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    hidden_states: list[torch.Tensor] | None = None
+    logits: torch.Tensor | None = None
     loss: torch.Tensor
-    extra_info: ModelForwardExtraLogInfo
+    extra_info: ModelForwardExtraLogInfo | None = None
+
+    def free_nongrad_feature(self):
+        """Release large intermediate tensors not needed for backward or
+        logging.
+
+        This method is called immediately after forward() in the micro-batch loop.
+        It releases large tensors (logits, hidden_states) while keeping:
+        - loss: needed for backward pass
+        - extra_info: lightweight logging info needed by post_micro_batch_forward()
+        """
+        self.hidden_states = None
+        self.logits = None
+
+    # TODO: Only for avoid BC. Should be removed later.
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    # TODO: Only for avoid BC. Should be removed later.
+    def __contains__(self, key):
+        return key in self.model_fields_set
 
 
 def _is_float8_available():

xtuner/v1/model/dense/dense.py

@@ -107,7 +107,7 @@ def forward(
         output["loss"] = loss
         output["logits"] = logits
         output["extra_info"] = extra_info
-        return ModelOutputs(**output)  # type: ignore[typeddict-item]
+        return ModelOutputs(**output)
 
     def build_embeddings(self, config: TransformerConfig):
         return nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)

xtuner/v1/model/moe/moe.py

@@ -79,10 +79,22 @@
 
 
 class MoEModelOutputs(ModelOutputs):
-    router_logits: NotRequired[dict[str, torch.Tensor]]
-    balancing_loss: NotRequired[torch.Tensor]
-    z_loss: NotRequired[torch.Tensor]
-    tokens_per_expert_global: NotRequired[torch.Tensor]
+    router_logits: dict[str, torch.Tensor] | None = None
+    balancing_loss: torch.Tensor | None = None
+    z_loss: torch.Tensor | None = None
+    tokens_per_expert_global: torch.Tensor
+
+    def free_nongrad_feature(self):
+        """Release large intermediate tensors not needed for backward or
+        logging.
+
+        This method is called immediately after forward() in the micro-batch loop.
+        It releases large tensors (logits, hidden_states) while keeping:
+        - loss: needed for backward pass
+        - extra_info: lightweight logging info needed by post_micro_batch_forward()
+        """
+        super().free_nongrad_feature()
+        self.router_logits = None
 
 
 class BalancingLossConfig(PydanticBaseModel):
@@ -482,7 +494,7 @@ def _micro_batch_forward(
 
             output["router_logits"] = router_logits_dict
 
-        return MoEModelOutputs(**output, logits=logits)  # type: ignore[typeddict-item]
+        return MoEModelOutputs(**output, logits=logits)
 
     def _forward(
         self,
@@ -583,7 +595,7 @@ def _forward(
         else:
             output["router_logits"] = None
 
-        return MoEModelOutputs(**output)  # type: ignore[typeddict-item]
+        return MoEModelOutputs(**output)
 
     def build_embeddings(self, config: MoEConfig):
         return nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)

xtuner/v1/utils/internal_metrics.py

@@ -198,7 +198,7 @@ def pop_metrics(self, data_batches: list[ModelItem]):
 
                 if (
                     self.internal_metrics_cfg.monitor_moe_load_balance_stats
-                    and (cur_tokens_per_expert := output.get("tokens_per_expert_global")) is not None
+                    and (cur_tokens_per_expert := output.tokens_per_expert_global) is not None
                 ):
                     # At this point, tokens_per_expert_global is already all-reduced into current rank.
                     # [num_layers, num_experts]
@@ -209,7 +209,7 @@ def pop_metrics(self, data_batches: list[ModelItem]):
 
                 if (
                     self.internal_metrics_cfg.monitor_moe_router_logits_stats
-                    and (cur_router_logits := output.get("router_logits")) is not None
+                    and (cur_router_logits := output.router_logits) is not None
                 ):
                     for layer_name, router_logits in cur_router_logits.items():
                         # [bsz, packed_len, num_experts]