[Enhance] Make ModelOutputs pydantic BaseModel

Author: HAOCHENYE

.github/workflows/claude.yml

@@ -9,6 +9,9 @@ on:
     types: [opened, assigned]
   pull_request_review:
     types: [submitted]
+  pull_request_target:
+      types: [opened, synchronize]
+      branches: [main]
 
 jobs:
   claude:
@@ -38,7 +41,7 @@ jobs:
           # Prompt A workaround for claude code action bug of `Fork` PR
           prompt: |
               REPO: ${{ github.repository }}
-              PR NUMBER: ${{ github.event.pull_request.number }}
+              PR NUMBER: ${{ github.event.pull_request.number || github.event.issue.number}}
 
               Please review this pull request.
 

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,25 @@ 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(self):
+        self.hidden_states = None
+        self.logits = None
+        self.extra_info = 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,14 @@
 
 
 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(self):
+        super().free()
+        self.router_logits = None
 
 
 class BalancingLossConfig(PydanticBaseModel):
@@ -482,7 +486,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 +587,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)