[Enhance] Make ModelOutputs pydantic BaseModel

ghstack-source-id: 067a5e1
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

@@ -12,6 +12,7 @@
 from typing_extensions import Self
 
 from xtuner.v1.loss.utils import sp_split
+from xtuner.v1.model.utils.misc import ModelForwardExtraLogInfo
 
 from .chunk_loss import ChunkLoss
 
@@ -195,6 +196,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)