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"]}