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/config/train_config.py

@@ -61,6 +61,7 @@ class TrainConfig(BaseConfig):
         val_split (Optional[float]): Ratio to split from training dataset for validation, e.g. 0.05. Defaults to None.
         val_step (int): Number of optimizer steps between validation runs. Defaults to 1000.
         neftune_alpha (float): NEFTune noise alpha, 0=disabled, typical: 5.0. Defaults to 0.0.
+        moe_aux_loss_coef (float): Weight applied to the MoE load-balancing loss. Defaults to 0.01.
         rollout_interval (int): Number of optimizer steps between online rollouts. Defaults to 512.
         rollout_temperature (float): Sampling temperature for online rollout. Defaults to 0.7.
         rollout_top_k (int): Top-k filtering for online rollout, 0=disable. Defaults to 0.
@@ -112,6 +113,7 @@ class TrainConfig(BaseConfig):
     val_split: Optional[float] = None
     val_step: int = 1000
     neftune_alpha: float = 0.0
+    moe_aux_loss_coef: float = 0.01
 
     rollout_interval: int = 512
     rollout_temperature: float = 0.7
@@ -187,7 +189,9 @@ class TrainConfig(BaseConfig):
             raise ValueError(f"rollout_top_p must be in (0, 1], got {v}")
         return v
 
-    @field_validator("rollout_top_k", "num_workers", "neftune_alpha")
+    @field_validator(
+        "rollout_top_k", "num_workers", "neftune_alpha", "moe_aux_loss_coef"
+    )
     def _validate_non_negative(cls, v):
         if v < 0:
             raise ValueError(f"must be non-negative, got {v}")

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

astrai/trainer/strategy.py

@@ -1,7 +1,7 @@
 """Training strategy implementations with factory pattern."""
 
 from abc import ABC, abstractmethod
-from typing import Callable, Dict, Union
+from typing import Callable, Dict, Optional, TypedDict, Union
 
 import torch
 import torch.nn as nn
@@ -13,6 +13,16 @@ from astrai.parallel.executor import broadcast_state_dict
 from astrai.trainer.rollout import RolloutResult
 
 
+class LossOutput(TypedDict):
+    loss: Tensor
+    metrics: Dict[str, Tensor]
+
+
+class LogprobsOutput(TypedDict):
+    logprobs: Tensor
+    aux_loss: Optional[Tensor]
+
+
 def move_to_device(batch: Dict[str, Tensor], device: str) -> Dict[str, Tensor]:
     """Move batch tensors to specified device with non-blocking transfer."""
     return {key: value.to(device, non_blocking=True) for key, value in batch.items()}
@@ -24,7 +34,7 @@ def get_logprobs(
     attn_mask: Tensor,
     loss_mask: Tensor,
     reduction: str,
-) -> Tensor:
+) -> LogprobsOutput:
     """Compute token-wise log probabilities from model outputs.
 
     Args:
@@ -46,10 +56,11 @@ def get_logprobs(
     shifted_input_ids = input_ids[:, 1:]
     shifted_loss_mask = loss_mask[:, 1:]
 
-    logits = model(
+    outputs = model(
         input_ids[:, :-1],
         attn_mask[:, :, :-1, :-1] if attn_mask.dim() == 4 else attn_mask[:, :-1],
-    )["logits"]
+    )
+    logits = outputs["logits"]
     log_probs = torch.log_softmax(logits.float(), dim=-1)
 
     token_logprobs = torch.gather(
@@ -57,13 +68,14 @@ def get_logprobs(
     ).squeeze(-1)
 
     if reduction == "mean":
-        return (token_logprobs * shifted_loss_mask).sum(dim=-1) / shifted_loss_mask.sum(
+        logprobs = (token_logprobs * shifted_loss_mask).sum(
             dim=-1
-        ).clamp(min=1.0)
+        ) / shifted_loss_mask.sum(dim=-1).clamp(min=1.0)
     elif reduction == "sum":
-        return (token_logprobs * shifted_loss_mask).sum(dim=-1)
+        logprobs = (token_logprobs * shifted_loss_mask).sum(dim=-1)
     else:
-        return token_logprobs * shifted_loss_mask
+        logprobs = token_logprobs * shifted_loss_mask
+    return {"logprobs": logprobs, "aux_loss": outputs.get("aux_loss")}
 
 
 def make_doc_boundary_mask(position_ids: Tensor) -> Tensor:
@@ -102,6 +114,7 @@ class BaseStrategy(ABC):
         self.model = model
         self.device = device
         self.executor = kwargs.pop("executor", None)
+        self.moe_aux_loss_coef = kwargs.pop("moe_aux_loss_coef", 0.01)
         self.extra_kwargs = kwargs
         self._rollout_runner = None
 
@@ -117,6 +130,33 @@ class BaseStrategy(ABC):
         """
         raise NotImplementedError
 
+    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
+        return self._normalize_output(self.compute_loss(batch))
+
+    def _loss_output(
+        self,
+        task_loss: Tensor,
+        metrics: Dict[str, Tensor],
+        aux_loss: Optional[Tensor] = None,
+    ) -> LossOutput:
+        total_loss = task_loss
+        if aux_loss is not None:
+            weighted_aux_loss = self.moe_aux_loss_coef * aux_loss
+            total_loss = total_loss + weighted_aux_loss
+            metrics["moe_aux_loss"] = aux_loss
+            metrics["moe_aux_loss_weighted"] = weighted_aux_loss
+        metrics["loss"] = total_loss
+        return {
+            "loss": total_loss,
+            "metrics": {name: value.detach() for name, value in metrics.items()},
+        }
+
+    @staticmethod
+    def _normalize_output(output: Union[LossOutput, Tensor]) -> LossOutput:
+        if isinstance(output, dict):
+            return output
+        return {"loss": output, "metrics": {"loss": output.detach()}}
+
     def supports_online(self) -> bool:
         """Whether this strategy can operate with a rollout runner.
 
@@ -153,17 +193,17 @@ class BaseStrategy(ABC):
         if self._rollout_runner is not None:
             self._rollout_runner.step()
 
-    def __call__(self, batch: Dict[str, Tensor]) -> Tensor:
+    def __call__(self, batch: Dict[str, Tensor]) -> LossOutput:
         """Run offline or online forward depending on runner injection."""
         if self._rollout_runner is None:
-            return self.compute_loss(batch)
+            return self.compute_loss_output(batch)
 
         result, is_fresh = self._rollout_runner(batch)
         if is_fresh:
             self._on_rollout_refresh()
 
         train_batch = self.prepare_from_rollout(result)
-        return self.compute_loss(train_batch)
+        return self.compute_loss_output(train_batch)
 
 
 class StrategyFactory(BaseFactory["BaseStrategy"]):
@@ -203,9 +243,13 @@ class SEQStrategy(BaseStrategy):
         self.label_smoothing = label_smoothing
 
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        return self.compute_loss_output(batch)["loss"]
+
+    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         batch = move_to_device(batch, self.device)
         input_ids, target_ids = batch["input_ids"], batch["target_ids"]
-        logits = self.model(input_ids=input_ids)["logits"]
+        outputs = self.model(input_ids=input_ids)
+        logits = outputs["logits"]
 
         loss = F.cross_entropy(
             input=logits.flatten(0, 1).float(),
@@ -213,7 +257,7 @@ class SEQStrategy(BaseStrategy):
             label_smoothing=self.label_smoothing,
         )
 
-        return loss
+        return self._loss_output(loss, {"task_loss": loss}, outputs.get("aux_loss"))
 
 
 @StrategyFactory.register("sft")
@@ -234,6 +278,9 @@ class SFTStrategy(BaseStrategy):
         self.label_smoothing = label_smoothing
 
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        return self.compute_loss_output(batch)["loss"]
+
+    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         batch = move_to_device(batch, self.device)
         input_ids, target_ids, position_ids, loss_mask = (
             batch["input_ids"],
@@ -245,9 +292,10 @@ class SFTStrategy(BaseStrategy):
         ignore_index = -100
         input_mask = make_doc_boundary_mask(position_ids)
         target_ids = target_ids.masked_fill(~loss_mask, ignore_index)
-        logits = self.model(
+        outputs = self.model(
             input_ids=input_ids, position_ids=position_ids, input_mask=input_mask
-        )["logits"]
+        )
+        logits = outputs["logits"]
 
         loss = F.cross_entropy(
             input=logits.flatten(0, 1).float(),
@@ -256,7 +304,7 @@ class SFTStrategy(BaseStrategy):
             label_smoothing=self.label_smoothing,
         )
 
-        return loss
+        return self._loss_output(loss, {"task_loss": loss}, outputs.get("aux_loss"))
 
 
 @StrategyFactory.register("dpo")
@@ -282,6 +330,9 @@ class DPOStrategy(BaseStrategy):
         self.reduction = reduction
 
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        return self.compute_loss_output(batch)["loss"]
+
+    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         batch = move_to_device(batch, self.device)
         chosen_ids, rejected_ids = batch["chosen"], batch["rejected"]
         chosen_mask, rejected_mask = batch["chosen_mask"], batch["rejected_mask"]
@@ -297,22 +348,25 @@ class DPOStrategy(BaseStrategy):
         )[None, None, :, :]  # [1, 1, S, S]
         full_mask = key_pad & causal  # [B*2, 1, S, S] — composed
 
-        log_pi = get_logprobs(
+        policy_output = get_logprobs(
             self.model,
             concat_ids,
             full_mask,
             concat_loss_mask,
             self.reduction,
         )
+        log_pi = policy_output["logprobs"]
+        aux_loss = policy_output["aux_loss"]
 
         with torch.no_grad():
-            log_ref = get_logprobs(
+            ref_output = get_logprobs(
                 self.ref_model,
                 concat_ids,
                 full_mask,
                 concat_loss_mask,
                 self.reduction,
             )
+            log_ref = ref_output["logprobs"]
 
         log_pi_chosen = log_pi[: chosen_ids.shape[0]]
         log_pi_rejected = log_pi[chosen_ids.shape[0] :]
@@ -325,7 +379,7 @@ class DPOStrategy(BaseStrategy):
         ratio_diff = pi_log_ratio - ref_log_ratio
         dpo_loss = -F.logsigmoid(self.beta * ratio_diff).mean()
 
-        return dpo_loss
+        return self._loss_output(dpo_loss, {"dpo_loss": dpo_loss}, aux_loss)
 
     def supports_online(self) -> bool:
         return True
@@ -398,6 +452,9 @@ class GRPOStrategy(BaseStrategy):
             self.old_model.load_state_dict(state_dict)
 
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        return self.compute_loss_output(batch)["loss"]
+
+    def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
         batch = move_to_device(batch, self.device)
         prompts = batch["prompts"]
         responses = batch["responses"]
@@ -438,16 +495,23 @@ class GRPOStrategy(BaseStrategy):
         # get_logprobs returns [B*G, S-1] (S = prompt_len + response_len).
         # Response token logprobs occupy the last ``response_len`` positions
         # (the first response token is predicted from the last prompt token).
-        token_log_probs_policy = get_logprobs(
+        policy_output = get_logprobs(
             self.model, full_sequences, attn_mask, full_masks, "none"
-        )[:, prompt_len - 1 :]
+        )
+        token_log_probs_policy = policy_output["logprobs"]
+        aux_loss = policy_output["aux_loss"]
+        token_log_probs_policy = token_log_probs_policy[:, prompt_len - 1 :]
         with torch.no_grad():
-            token_log_probs_old = get_logprobs(
+            old_output = get_logprobs(
                 self.old_model, full_sequences, attn_mask, full_masks, "none"
-            )[:, prompt_len - 1 :]
-            token_log_probs_ref = get_logprobs(
+            )
+            token_log_probs_old = old_output["logprobs"]
+            token_log_probs_old = token_log_probs_old[:, prompt_len - 1 :]
+            ref_output = get_logprobs(
                 self.ref_model, full_sequences, attn_mask, full_masks, "none"
-            )[:, prompt_len - 1 :]
+            )
+            token_log_probs_ref = ref_output["logprobs"]
+            token_log_probs_ref = token_log_probs_ref[:, prompt_len - 1 :]
 
         # Reshape to [B, G, response_len]
         token_log_probs_policy = token_log_probs_policy.view(batch_size, group_size, -1)
@@ -480,9 +544,12 @@ class GRPOStrategy(BaseStrategy):
         kl_per_token = r - torch.log(r + eps) - 1.0
         kl_penalty = self.kl_coef * (kl_per_token * token_masks).sum() / token_count
 
-        total_loss = policy_loss + kl_penalty
-
-        return total_loss
+        task_loss = policy_loss + kl_penalty
+        return self._loss_output(
+            task_loss,
+            {"policy_loss": policy_loss, "kl_loss": kl_penalty},
+            aux_loss,
+        )
 
     def supports_online(self) -> bool:
         return True

astrai/trainer/train_callback.py

@@ -260,11 +260,28 @@ class MetricCallback(TrainCallback):
         }
 
     def _metrics(self, context: TrainContext, names):
-        return {
-            m: self._metric_funcs[m](context)
-            for m in names
-            if self._metric_funcs[m](context) is not None
-        }
+        metrics = dict(context.metrics)
+        for name in names:
+            metric_fn = self._metric_funcs.get(name)
+            if metric_fn is None:
+                continue
+            value = metric_fn(context)
+            if value is not None:
+                metrics[name] = value
+        selected = set(context.metrics) | set(names)
+        selected.discard("*")
+        result = {name: metrics[name] for name in selected if name in metrics}
+        if context.world_size > 1 and dist.is_initialized() and result:
+            metric_names = sorted(result)
+            values = torch.tensor(
+                [result[name] for name in metric_names],
+                dtype=torch.float32,
+                device=get_current_device(),
+            )
+            dist.all_reduce(values, op=dist.ReduceOp.SUM)
+            values /= context.world_size
+            result.update(zip(metric_names, values.tolist()))
+        return result
 
     @only_on_rank(0)
     def _append(self, event_type: str, context: TrainContext, **extra):
@@ -286,8 +303,8 @@ class MetricCallback(TrainCallback):
 
         with torch.no_grad():
             for batch in context.val_dataloader:
-                loss = context.strategy(batch)
-                total_loss += loss.item()
+                loss_output = context.strategy(batch)
+                total_loss += loss_output["loss"].item()
                 num_batches += 1
 
         if context.world_size > 1 and dist.is_initialized():

astrai/trainer/train_context.py

@@ -38,6 +38,7 @@ class TrainContext:
     epoch: int = field(default=0)
     consumed_samples: int = field(default=0)
     loss: float = field(default=0.0)
+    metrics: Dict[str, float] = field(default_factory=dict)
     grad_norm: Optional[float] = field(default=None)
     grad_snr_tracker: GradSNRTracker = field(default_factory=GradSNRTracker)
     val_dataloader: Optional[DataLoader] = field(default=None)
@@ -221,6 +222,7 @@ class TrainContextBuilder:
                         obj.load_state_dict(extra[name])
 
         strategy_kwargs = dict(cfg.extra_kwargs)
+        strategy_kwargs.setdefault("moe_aux_loss_coef", cfg.moe_aux_loss_coef)
 
         needs_ref = cfg.strategy in (
             "dpo",

astrai/trainer/trainer.py

@@ -82,9 +82,13 @@ class Trainer:
                         break
                     with executor.accumulate(context.model):
                         self._call_callbacks("on_batch_begin", context)
-                        loss = context.strategy(batch)
-                        context.loss = loss.item()
-                        stand_loss = loss / executor.grad_accum_steps
+                        loss_output = context.strategy(batch)
+                        context.loss = loss_output["loss"].item()
+                        context.metrics = {
+                            name: value.item()
+                            for name, value in loss_output["metrics"].items()
+                        }
+                        stand_loss = loss_output["loss"] / executor.grad_accum_steps
                         executor.backward(stand_loss)
                         context.consumed_samples += (
                             context.config.batch_per_device * context.world_size