feat: add MoE auxiliary loss metrics

- Propagates MoE load-balancing loss through model outputs
- Logs task, auxiliary, and weighted losses across strategies
- Computes only explicitly requested callback metrics
- Preserves tensor compute_loss API and adds regression tests

astrai/model/components/decoder_block.py

@@ -1,5 +1,5 @@
 from dataclasses import asdict
-from typing import Optional
+from typing import Optional, TypedDict
 
 import torch.nn as nn
 from torch import Tensor
@@ -10,6 +10,11 @@ from astrai.model.components.mlp import FFNFactory
 from astrai.model.components.norm import RMSNorm
 
 
+class DecoderOutput(TypedDict):
+    hidden_states: Tensor
+    aux_loss: Optional[Tensor]
+
+
 class DecoderBlock(nn.Module):
     def __init__(self, config, layer_id: int):
         super().__init__()
@@ -48,7 +53,7 @@ class DecoderBlock(nn.Module):
         attention_mask: Optional[Tensor] = None,
         kv_cache: Optional[KVCache] = None,
         is_causal: bool = False,
-    ) -> Tensor:
+    ) -> DecoderOutput:
         attn_output = self.attention(
             self.input_norm(x),
             rotary_emb,
@@ -57,6 +62,8 @@ class DecoderBlock(nn.Module):
             is_causal,
         )
         x = attn_output + x
-        x = self.mlp(self.post_attention_norm(x)) + x
+        normalized = self.post_attention_norm(x)
+        mlp_output = self.mlp(normalized)
+        x = mlp_output["hidden_states"] + x
 
-        return x
+        return {"hidden_states": x, "aux_loss": mlp_output["aux_loss"]}

astrai/model/components/mlp.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, TypedDict
 
 import torch
 import torch.nn as nn
@@ -13,6 +13,16 @@ class FFNFactory(BaseFactory[nn.Module]):
     pass
 
 
+class FFNOutput(TypedDict):
+    hidden_states: Tensor
+    aux_loss: Optional[Tensor]
+
+
+class RoutedOutput(TypedDict):
+    hidden_states: Tensor
+    aux_loss: Optional[Tensor]
+
+
 @FFNFactory.register("mlp")
 class MLP(nn.Module):
     def __init__(self, dim: int, dim_ffn: int, down_init_std: float = 0.02):
@@ -21,10 +31,10 @@ class MLP(nn.Module):
         self.gate = Linear(dim, dim_ffn)
         self.down = Linear(dim_ffn, dim, init_std=down_init_std)
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(self, x: Tensor) -> FFNOutput:
         gated = self.up(x) * F.silu(self.gate(x))
         out = self.down(gated)
-        return out
+        return {"hidden_states": out, "aux_loss": None}
 
 
 @FFNFactory.register("moe")
@@ -76,22 +86,26 @@ class DeepSeekMoE(nn.Module):
             ]
         )
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(self, x: Tensor) -> FFNOutput:
+        include_aux_loss = self.training and torch.is_grad_enabled()
         bsz, seq_len, dim = x.shape
         x_flat = x.view(-1, dim)
 
         shared_out = self._shared_forward(x_flat)
-        routed_out = self._routed_forward(x_flat)
+        routed_output = self._routed_forward(x_flat, include_aux_loss)
 
-        out = (shared_out + routed_out).view(bsz, seq_len, dim)
-        return out
+        out = (shared_out + routed_output["hidden_states"]).view(bsz, seq_len, dim)
+        return {"hidden_states": out, "aux_loss": routed_output["aux_loss"]}
 
     def _shared_forward(self, x: Tensor) -> Tensor:
         if self.n_shared_experts == 0:
             return torch.zeros_like(x)
-        return sum(e(x) for e in self.shared_experts) / self.n_shared_experts
+        return (
+            sum(e(x)["hidden_states"] for e in self.shared_experts)
+            / self.n_shared_experts
+        )
 
-    def _routed_forward(self, x: Tensor) -> Tensor:
+    def _routed_forward(self, x: Tensor, include_aux_loss: bool) -> RoutedOutput:
         N, D = x.shape
         K = self.n_activated_experts
 
@@ -102,15 +116,26 @@ class DeepSeekMoE(nn.Module):
         if self.norm_topk_prob:
             topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
 
+        aux_loss = None
+        if include_aux_loss:
+            expert_load = F.one_hot(
+                topk_indices, num_classes=self.n_routed_experts
+            ).float()
+            expert_load = expert_load.mean(dim=(0, 1))
+            router_prob = router_probs.float().mean(dim=0)
+            aux_loss = self.n_routed_experts * (expert_load * router_prob).sum()
+
         output = torch.zeros(N, D, device=x.device, dtype=x.dtype)
         for expert_idx in range(self.n_routed_experts):
             expert_mask = topk_indices == expert_idx
             token_idx, k_idx = expert_mask.nonzero(as_tuple=True)
             if token_idx.numel() == 0:
                 continue
+            expert = self.routed_experts[expert_idx]
             expert_input = x[token_idx]
-            expert_output = self.routed_experts[expert_idx](expert_input)
+            expert_output = expert(expert_input)["hidden_states"]
+
             weights = topk_weights[token_idx, k_idx].unsqueeze(-1)
             output.index_add_(0, token_idx, expert_output * weights)
 
-        return output
+        return {"hidden_states": output, "aux_loss": aux_loss}

astrai/model/encoder.py

@@ -70,7 +70,7 @@ class EmbeddingEncoder(AutoModel):
         attn_mask = process_attention_mask(input_mask)
 
         for layer in self.layers:
-            x = layer(x, rotary_emb, attn_mask)
+            x = layer(x, rotary_emb, attn_mask)["hidden_states"]
 
         hidden_states = self.norm(x)
 

astrai/model/transformer.py

@@ -113,10 +113,23 @@ class AutoRegressiveLM(AutoModel):
         attn_mask = process_attention_mask(input_mask)
         use_sdpa_causal_mask = attn_mask is None
 
+        aux_losses = []
         for layer in self.layers:
-            x = layer(x, rotary_emb, attn_mask, kv_cache, use_sdpa_causal_mask)
+            layer_output = layer(
+                x,
+                rotary_emb,
+                attn_mask,
+                kv_cache,
+                use_sdpa_causal_mask,
+            )
+            x = layer_output["hidden_states"]
+            if layer_output["aux_loss"] is not None:
+                aux_losses.append(layer_output["aux_loss"])
 
         hidden_states = self.norm(x)
         logits = self.lm_head(hidden_states)
 
-        return {"logits": logits, "hidden_states": hidden_states}
+        output = {"logits": logits, "hidden_states": hidden_states}
+        if aux_losses:
+            output["aux_loss"] = torch.stack(aux_losses).mean()
+        return output