feat: add moe auxloss and metrics

astrai/model/components/mlp.py

@@ -1,4 +1,4 @@
-from typing import Optional, TypedDict
+from typing import List, Optional, TypedDict
 
 import torch
 import torch.nn as nn
@@ -70,6 +70,7 @@ class DeepSeekMoE(nn.Module):
         )
 
         self.router = Linear(dim, n_routed_experts, bias=False)
+        self._router_probs: Optional[Tensor] = None
         moe_scale = 1 / max(n_shared_experts, 1) + 1 / n_activated_experts
         down_init_std = 0.02 / (2 * n_layers * moe_scale) ** 0.5
 
@@ -111,6 +112,7 @@ class DeepSeekMoE(nn.Module):
 
         router_logits = self.router(x)
         router_probs = torch.softmax(router_logits.float(), dim=-1).to(x.dtype)
+        self._router_probs = router_probs.detach()
 
         topk_weights, topk_indices = torch.topk(router_probs, K, dim=-1)
         if self.norm_topk_prob:
@@ -139,3 +141,12 @@ class DeepSeekMoE(nn.Module):
             output.index_add_(0, token_idx, expert_output * weights)
 
         return {"hidden_states": output, "aux_loss": aux_loss}
+
+    @staticmethod
+    def collect_router_probs(module: nn.Module) -> List[Tensor]:
+        """Recursively collect router_probs from all DeepSeekMoE submodules."""
+        probs: List[Tensor] = []
+        for m in module.modules():
+            if isinstance(m, DeepSeekMoE) and m._router_probs is not None:
+                probs.append(m._router_probs)
+        return probs

astrai/model/transformer.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Mapping, Optional
+from typing import Any, Dict, List, Mapping, Optional
 
 import torch
 import torch.nn as nn
@@ -10,6 +10,7 @@ from astrai.model.automodel import AutoModel, ModelFactory
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding
 from astrai.model.components.linear import Linear
+from astrai.model.components.mlp import DeepSeekMoE
 from astrai.model.components.norm import RMSNorm
 from astrai.model.components.rope import RotaryEmbedding
 
@@ -133,3 +134,7 @@ class AutoRegressiveLM(AutoModel):
         if aux_losses:
             output["aux_loss"] = torch.stack(aux_losses).mean()
         return output
+
+    def get_moe_router_probs(self) -> List[Tensor]:
+        """Return router_probs from all MoE layers for strategy-side aux loss."""
+        return DeepSeekMoE.collect_router_probs(self)