feat: add moe auxloss and metrics

astrai/config/model_config.py

@@ -97,6 +97,7 @@ class AutoRegressiveLMConfig(BaseModelConfig):
     norm_topk_prob: bool = True
     decoder_sparse_step: int = 1
     mlp_only_layers: Optional[list[int]] = None
+    moe_aux_loss_coef: float = 0.01
 
     @field_validator("attn_type")
     def _validate_attn_type(cls, v: str) -> str:

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)

astrai/trainer/metric_util.py

@@ -88,3 +88,23 @@ def ctx_get_grad_snr(ctx):
     if tracker is None:
         return None
     return tracker.snr
+
+
+def ctx_get_moe_aux_loss(ctx):
+    return ctx.strategy._moe_metrics.get("aux_loss")
+
+
+def ctx_get_router_entropy(ctx):
+    return ctx.strategy._moe_metrics.get("router_entropy")
+
+
+def ctx_get_dead_expert_fraction(ctx):
+    return ctx.strategy._moe_metrics.get("dead_expert_fraction")
+
+
+def ctx_get_load_imbalance_mean(ctx):
+    return ctx.strategy._moe_metrics.get("load_imbalance_mean")
+
+
+def ctx_get_load_imbalance_max(ctx):
+    return ctx.strategy._moe_metrics.get("load_imbalance_max")

astrai/trainer/strategy.py

@@ -1,7 +1,7 @@
 """Training strategy implementations with factory pattern."""
 
 from abc import ABC, abstractmethod
-from typing import Callable, Dict, Optional, TypedDict, Union
+from typing import Callable, Dict, List, Optional, TypedDict, Union
 
 import torch
 import torch.nn as nn
@@ -94,6 +94,97 @@ def make_doc_boundary_mask(position_ids: Tensor) -> Tensor:
     return (same_doc & causal).unsqueeze(1)
 
 
+def _load_balancing_loss(router_probs: Tensor) -> Tensor:
+    """Compute MoE load balancing auxiliary loss from router probabilities.
+
+    Implements the Switch Transformer load balancing loss (eq. 4-6).
+    Encourages tokens to be uniformly distributed across experts.
+
+    Args:
+        router_probs: (N, num_experts) tensor of softmax router probabilities.
+
+    Returns:
+        Scalar aux loss = num_experts * sum(f_i * P_i).
+    """
+    num_experts = router_probs.size(-1)
+    # f_i: fraction of tokens dispatched to expert i (soft mean)
+    f_i = router_probs.mean(dim=0)
+    # P_i: average routing probability for expert i
+    P_i = router_probs.mean(dim=0)
+    return num_experts * torch.sum(f_i * P_i)
+
+
+def _collect_moe_diagnostics(
+    router_probs_list: List[Tensor],
+    top_k: int,
+) -> Dict[str, float]:
+    """Collect MoE routing diagnostic metrics from router probabilities.
+
+    Args:
+        router_probs_list: List of (N, num_experts) router probability tensors,
+            one per MoE layer.
+        top_k: Number of top experts selected per token.
+
+    Returns:
+        Dict with keys: router_entropy, dead_expert_fraction,
+        load_imbalance_mean, load_imbalance_max.  Values are averaged
+        across layers.
+    """
+    layer_entropies: List[Tensor] = []
+    layer_dead_fractions: List[Tensor] = []
+    layer_imbalance_means: List[Tensor] = []
+    layer_imbalance_maxs: List[Tensor] = []
+
+    for probs in router_probs_list:
+        probs = probs.detach().to(dtype=torch.float32)
+        if probs.ndim == 0 or probs.shape[-1] == 0:
+            continue
+        probs = probs.reshape(-1, probs.shape[-1])
+        if probs.numel() == 0:
+            continue
+
+        num_experts = probs.shape[-1]
+        num_tokens = probs.shape[0]
+
+        # Router entropy
+        entropy = -(probs * torch.log(probs.clamp_min(1e-8))).sum(dim=-1).mean()
+
+        # Top-k expert selection
+        selected_experts = torch.topk(probs, top_k, dim=-1).indices  # [tokens, top_k]
+        expert_mask = F.one_hot(selected_experts, num_experts)  # [tokens, top_k, E]
+        expert_counts = expert_mask.sum(dim=(0, 1)).to(dtype=torch.float32)  # [E]
+
+        # Ideal load: tokens * top_k / num_experts
+        ideal_load = (num_tokens * top_k) / max(num_experts, 1)
+
+        # Load imbalance ratios
+        load_ratios = expert_counts / max(ideal_load, 1.0)
+        imbalance_mean = (load_ratios - 1.0).abs().mean()
+        imbalance_max = load_ratios.max()
+        dead_fraction = (expert_counts == 0).to(dtype=torch.float32).mean()
+
+        layer_entropies.append(entropy)
+        layer_dead_fractions.append(dead_fraction)
+        layer_imbalance_means.append(imbalance_mean)
+        layer_imbalance_maxs.append(imbalance_max)
+
+    if not layer_entropies:
+        return {}
+
+    return {
+        "router_entropy": float(torch.stack(layer_entropies).mean().cpu().item()),
+        "dead_expert_fraction": float(
+            torch.stack(layer_dead_fractions).mean().cpu().item()
+        ),
+        "load_imbalance_mean": float(
+            torch.stack(layer_imbalance_means).mean().cpu().item()
+        ),
+        "load_imbalance_max": float(
+            torch.stack(layer_imbalance_maxs).mean().cpu().item()
+        ),
+    }
+
+
 class BaseStrategy(ABC):
     """Abstract base class for training strategies.
 
@@ -115,6 +206,7 @@ class BaseStrategy(ABC):
         self.device = device
         self.executor = kwargs.pop("executor", None)
         self.moe_aux_loss_coef = kwargs.pop("moe_aux_loss_coef", 0.01)
+        self._moe_metrics: Dict[str, float] = {}
         self.extra_kwargs = kwargs
         self._rollout_runner = None
 
@@ -145,6 +237,7 @@ class BaseStrategy(ABC):
             total_loss = total_loss + weighted_aux_loss
             metrics["moe_aux_loss"] = aux_loss
             metrics["moe_aux_loss_weighted"] = weighted_aux_loss
+            self._refresh_moe_diagnostics(aux_loss)
         metrics["loss"] = total_loss
         return {
             "loss": total_loss,
@@ -188,6 +281,23 @@ class BaseStrategy(ABC):
         """
         pass
 
+    def _refresh_moe_diagnostics(self, aux_loss: Tensor) -> None:
+        """Collect MoE routing diagnostics from model router probs.
+
+        Populates ``self._moe_metrics`` with router entropy, dead expert
+        fraction, load imbalance, and aux_loss.  Called from
+        :meth:`_loss_output` when an MoE aux loss is present.
+        """
+        router_probs_list: List[Tensor] = self.model.get_moe_router_probs()
+        if not router_probs_list:
+            self._moe_metrics = {}
+            return
+        self._moe_metrics = _collect_moe_diagnostics(
+            router_probs_list,
+            self.model.config.n_activated_experts,
+        )
+        self._moe_metrics["aux_loss"] = float(aux_loss.detach().cpu().item())
+
     def on_optimizer_step(self):
         """Advance online rollout state after a successful optimizer step."""
         if self._rollout_runner is not None:
@@ -230,6 +340,7 @@ class SEQStrategy(BaseStrategy):
     """Standard next-token prediction training strategy.
 
     Computes cross-entropy loss for next token prediction.
+    Optionally adds MoE load balancing auxiliary loss.
     """
 
     def __init__(
@@ -265,6 +376,7 @@ class SFTStrategy(BaseStrategy):
     """Supervised Fine-tuning strategy with loss masking.
 
     Applies cross-entropy loss only to tokens where loss_mask is True.
+    Optionally adds MoE load balancing auxiliary loss.
     """
 
     def __init__(

astrai/trainer/train_callback.py

@@ -17,10 +17,15 @@ from astrai.parallel import only_on_rank
 from astrai.parallel.setup import get_current_device
 from astrai.serialization import Checkpoint
 from astrai.trainer.metric_util import (
+    ctx_get_dead_expert_fraction,
     ctx_get_grad_norm,
     ctx_get_grad_snr,
+    ctx_get_load_imbalance_max,
+    ctx_get_load_imbalance_mean,
     ctx_get_loss,
     ctx_get_lr,
+    ctx_get_moe_aux_loss,
+    ctx_get_router_entropy,
     ctx_get_val_loss,
 )
 from astrai.trainer.train_context import TrainContext
@@ -257,6 +262,11 @@ class MetricCallback(TrainCallback):
             "val_loss": ctx_get_val_loss,
             "grad_norm": ctx_get_grad_norm,
             "grad_snr": ctx_get_grad_snr,
+            "moe_aux_loss": ctx_get_moe_aux_loss,
+            "router_entropy": ctx_get_router_entropy,
+            "dead_expert_fraction": ctx_get_dead_expert_fraction,
+            "load_imbalance_mean": ctx_get_load_imbalance_mean,
+            "load_imbalance_max": ctx_get_load_imbalance_max,
         }
 
     def _metrics(self, context: TrainContext, names):

scripts/tools/train.py

@@ -289,6 +289,13 @@ _START_METHODS = ["spawn", "fork", "forkserver"]
     group="Data Loading",
     help="Label smoothing.",
 )
+@opt(
+    "--moe_aux_loss_coef",
+    type=float,
+    default=0.01,
+    group="Algorithm",
+    help="MoE load balancing auxiliary loss coefficient (0=disable).",
+)
 @opt(
     "--rollout_interval",
     type=int,
@@ -813,6 +820,7 @@ def train(
         rollout_top_p=rollout_top_p,
         rollout_max_tokens=rollout_max_tokens,
         reward_model_fn=reward_model_fn,
+        moe_aux_loss_coef=kwargs.pop("moe_aux_loss_coef", 0.01),
     )
 
     trainer = Trainer(train_config)

tests/module/test_forward_configs.py

@@ -265,6 +265,103 @@ def test_moe_defaults_preserve_normalized_routing():
     assert model.layers[0].mlp.norm_topk_prob is True
 
 
+def test_moe_router_probs_populated_after_forward():
+    """Verify DeepSeekMoE._router_probs is set after forward in training mode."""
+    from astrai.config.model_config import AutoRegressiveLMConfig
+    from astrai.model.components.mlp import DeepSeekMoE
+
+    config = AutoRegressiveLMConfig(
+        **TINY_CONFIG,
+        ffn_type="moe",
+        n_routed_experts=4,
+        n_shared_experts=1,
+        n_activated_experts=2,
+        topk_method="greedy",
+    )
+    model = AutoRegressiveLM(config)
+    model.train()
+    input_ids = torch.randint(0, config.vocab_size, (2, 8))
+
+    with torch.enable_grad():
+        model(input_ids)
+
+    # All MoE layers should have router_probs set
+    moe_layers = [m for m in model.modules() if isinstance(m, DeepSeekMoE)]
+    assert len(moe_layers) > 0
+    for layer in moe_layers:
+        assert layer._router_probs is not None
+        assert layer._router_probs.ndim == 2
+        assert layer._router_probs.shape[-1] == 4  # n_routed_experts
+
+
+def test_get_moe_router_probs_moe_model():
+    """Verify get_moe_router_probs() returns a list of tensors for MoE models."""
+    from astrai.config.model_config import AutoRegressiveLMConfig
+
+    config = AutoRegressiveLMConfig(
+        **TINY_CONFIG,
+        ffn_type="moe",
+        n_routed_experts=4,
+        n_shared_experts=1,
+        n_activated_experts=2,
+    )
+    model = AutoRegressiveLM(config)
+    model.train()
+
+    with torch.enable_grad():
+        model(torch.randint(0, config.vocab_size, (2, 8)))
+
+    probs = model.get_moe_router_probs()
+    assert isinstance(probs, list)
+    assert len(probs) == 2  # num_hidden_layers
+    for p in probs:
+        assert p.ndim == 2
+        assert p.shape[-1] == 4
+
+
+def test_get_moe_router_probs_non_moe_model():
+    """Verify get_moe_router_probs() returns empty list for non-MoE models."""
+    from astrai.config.model_config import AutoRegressiveLMConfig
+
+    config = AutoRegressiveLMConfig(**TINY_CONFIG, ffn_type="mlp")
+    model = AutoRegressiveLM(config)
+
+    probs_untrained = model.get_moe_router_probs()
+    assert probs_untrained == []
+
+    model.train()
+    with torch.enable_grad():
+        model(torch.randint(0, config.vocab_size, (2, 8)))
+
+    probs = model.get_moe_router_probs()
+    assert probs == []
+
+
+def test_collect_router_probs_static_method():
+    """Verify DeepSeekMoE.collect_router_probs static method."""
+    from astrai.model.components.mlp import DeepSeekMoE
+
+    moe = DeepSeekMoE(
+        dim=8,
+        dim_ffn=16,
+        n_routed_experts=4,
+        n_shared_experts=1,
+        n_activated_experts=2,
+    )
+    moe.train()
+    with torch.enable_grad():
+        moe(torch.randn(2, 8, 8))
+
+    # collect_router_probs should find the MoE layer
+    probs = DeepSeekMoE.collect_router_probs(moe)
+    assert len(probs) == 1
+    assert probs[0].shape[-1] == 4
+
+    # On a plain MLP module, should return empty
+    mlp_module = MLP(8, 16)
+    assert DeepSeekMoE.collect_router_probs(mlp_module) == []
+
+
 def test_moe_aux_loss_only_emitted_during_training():
     from astrai.config.model_config import AutoRegressiveLMConfig
 

tests/module/test_model_forward.py

@@ -0,0 +1,316 @@
+"""Smoke tests for MoE aux loss and diagnostic metrics integration.
+
+Does NOT load real data or weights.  Uses a tiny randomly-initialized
+MoE model and verifies that aux loss computation and MoE routing
+diagnostics flow end‑to‑end through the strategy layer.
+"""
+
+import pytest
+import torch
+
+from astrai.config.model_config import AutoRegressiveLMConfig
+from astrai.model.components.mlp import DeepSeekMoE
+from astrai.model.transformer import AutoRegressiveLM
+from astrai.trainer.strategy import (
+    SEQStrategy,
+    SFTStrategy,
+    StrategyFactory,
+    _collect_moe_diagnostics,
+    _load_balancing_loss,
+)
+from tests.helpers import TINY_CONFIG
+
+# ── helpers ──────────────────────────────────────────────────────────
+
+
+def _make_tiny_moe_config(**overrides) -> AutoRegressiveLMConfig:
+    return AutoRegressiveLMConfig(
+        **{
+            **TINY_CONFIG,
+            "ffn_type": "moe",
+            "n_routed_experts": 4,
+            "n_shared_experts": 1,
+            "n_activated_experts": 2,
+            "topk_method": "greedy",
+            **overrides,
+        }
+    )
+
+
+def _make_model(config=None) -> AutoRegressiveLM:
+    if config is None:
+        config = _make_tiny_moe_config()
+    return AutoRegressiveLM(config)
+
+
+# ── _collect_moe_diagnostics unit tests ─────────────────────────────
+
+
+def test_collect_moe_diagnostics_returns_all_keys():
+    """_collect_moe_diagnostics should return the four expected keys."""
+    # Simulate two MoE layers with uniform routing probabilities
+    probs = torch.ones(128, 4) / 4.0
+    diag = _collect_moe_diagnostics([probs, probs], top_k=2)
+
+    assert set(diag.keys()) == {
+        "router_entropy",
+        "dead_expert_fraction",
+        "load_imbalance_mean",
+        "load_imbalance_max",
+    }
+    for v in diag.values():
+        assert isinstance(v, float)
+
+
+def test_collect_moe_diagnostics_empty_list():
+    """Empty list returns empty dict."""
+    assert _collect_moe_diagnostics([], top_k=2) == {}
+
+
+def test_collect_moe_diagnostics_uniform_routing():
+    """Uniform routing probabilities with top_k=2 → tie-breaking by index.
+
+    torch.topk breaks ties by index, so with equal probabilities
+    experts 0 and 1 always win over experts 2 and 3:
+      - dead_expert_fraction = 2/4 = 0.5
+      - load_ratios = [2, 2, 0, 0] → |ratio-1| = [1, 1, 1, 1] → mean = 1.0
+      - load_imbalance_max = 2.0
+    """
+    probs = torch.ones(128, 4) / 4.0
+    diag = _collect_moe_diagnostics([probs], top_k=2)
+
+    assert diag["dead_expert_fraction"] == pytest.approx(0.5, abs=1e-6)
+    assert diag["load_imbalance_mean"] == pytest.approx(1.0, abs=1e-6)
+    assert diag["load_imbalance_max"] == pytest.approx(2.0, abs=1e-6)
+
+
+def test_collect_moe_diagnostics_max_entropy():
+    """Uniform probabilities should give log(num_experts) entropy."""
+    num_experts = 4
+    probs = torch.ones(128, num_experts) / num_experts
+    diag = _collect_moe_diagnostics([probs], top_k=2)
+    expected_entropy = float(torch.log(torch.tensor(num_experts, dtype=torch.float32)))
+    assert diag["router_entropy"] == pytest.approx(expected_entropy, abs=1e-5)
+
+
+# ── _load_balancing_loss unit tests ──────────────────────────────────
+
+
+def test_load_balancing_loss_shape_and_range():
+    """Verify _load_balancing_loss returns a non-negative scalar tensor."""
+    probs = torch.randn(64, 8).softmax(dim=-1)
+    loss = _load_balancing_loss(probs)
+    assert loss.ndim == 0
+    assert loss.item() >= 0
+
+
+def test_load_balancing_loss_uniform_minimum():
+    """Uniform routing gives the lowest possible load balancing loss."""
+    probs = torch.ones(64, 8) / 8.0
+    loss = _load_balancing_loss(probs).item()
+
+    # Very skewed routing should give higher loss
+    skewed = torch.zeros(64, 8)
+    skewed[:, 0] = 1.0
+    skewed[:, 1] = 1.0
+    skewed = skewed / skewed.sum(dim=-1, keepdim=True)
+    skewed_loss = _load_balancing_loss(skewed).item()
+
+    assert loss < skewed_loss
+
+
+# ── SEQStrategy integration tests ────────────────────────────────────
+
+
+class TestSEQStrategyMoE:
+    """End‑to‑end tests for SEQStrategy with MoE aux loss."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self, device):
+        self.device = device
+        self.config = _make_tiny_moe_config()
+        self.model = _make_model(self.config).to(device)
+        self.model.train()
+
+    def _make_batch(self, batch_size=2, seq_len=8):
+        vocab = self.config.vocab_size
+        input_ids = torch.randint(0, vocab, (batch_size, seq_len))
+        # target = input shifted right
+        target_ids = torch.randint(0, vocab, (batch_size, seq_len))
+        return {"input_ids": input_ids, "target_ids": target_ids}
+
+    def test_compute_loss_returns_scalar(self):
+        """compute_loss should return a scalar tensor."""
+        strategy = SEQStrategy(
+            self.model,
+            self.device,
+            moe_aux_loss_coef=0.01,
+        )
+        loss = strategy.compute_loss(self._make_batch())
+        assert loss.ndim == 0
+        assert loss.requires_grad
+
+    def test_compute_loss_output_has_metrics(self):
+        """compute_loss_output dict with moe_aux_loss_coef > 0 includes MoE metrics."""
+        strategy = SEQStrategy(
+            self.model,
+            self.device,
+            moe_aux_loss_coef=0.01,
+        )
+        output = strategy.compute_loss_output(self._make_batch())
+
+        assert "loss" in output
+        assert "metrics" in output
+        assert output["loss"].ndim == 0
+        assert output["loss"].requires_grad
+
+        metrics = output["metrics"]
+        # MoE metrics should appear when coef > 0 and model has MoE layers
+        for key in ("moe_aux_loss", "moe_aux_loss_weighted", "task_loss", "loss"):
+            assert key in metrics, f"Missing metric: {key}"
+            assert isinstance(metrics[key], float)
+
+    def test_moe_metrics_populated_after_forward(self):
+        """strategy._moe_metrics populated after compute_loss_output."""
+        strategy = SEQStrategy(
+            self.model,
+            self.device,
+            moe_aux_loss_coef=0.01,
+        )
+        strategy.compute_loss_output(self._make_batch())
+
+        moe_metrics = strategy._moe_metrics
+        assert moe_metrics, "_moe_metrics should not be empty for MoE model"
+        for key in (
+            "aux_loss",
+            "router_entropy",
+            "dead_expert_fraction",
+            "load_imbalance_mean",
+            "load_imbalance_max",
+        ):
+            assert key in moe_metrics, f"Missing _moe_metrics key: {key}"
+            assert isinstance(moe_metrics[key], float)
+
+    def test_zero_coef_zeroes_weighted_aux(self):
+        """moe_aux_loss_coef=0 → weighted_aux_loss is zero, task_loss == loss."""
+        strategy = SEQStrategy(
+            self.model,
+            self.device,
+            moe_aux_loss_coef=0.0,
+        )
+        output = strategy.compute_loss_output(self._make_batch())
+        metrics = output["metrics"]
+
+        # task_loss and loss should be equal (aux weighted by zero)
+        assert "task_loss" in metrics
+        assert "loss" in metrics
+        assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
+
+        # weighted aux loss is zero
+        assert metrics.get("moe_aux_loss_weighted") == pytest.approx(0.0, abs=1e-6)
+
+        # MoE diagnostics are still collected (monitoring purposes)
+        assert strategy._moe_metrics
+        assert "router_entropy" in strategy._moe_metrics
+
+    def test_aux_loss_added_to_total_loss(self):
+        """Total loss > task_loss when moe_aux_loss_coef > 0."""
+        strategy = SEQStrategy(
+            self.model,
+            self.device,
+            moe_aux_loss_coef=0.01,
+        )
+        output = strategy.compute_loss_output(self._make_batch())
+        assert output["metrics"]["loss"] > output["metrics"]["task_loss"] + 1e-12
+
+    def test_factory_creates_strategy_with_coef(self):
+        """StrategyFactory.create passes moe_aux_loss_coef to strategy."""
+        strategy = StrategyFactory.create(
+            "seq",
+            model=self.model,
+            device=self.device,
+            moe_aux_loss_coef=0.02,
+        )
+        assert strategy.moe_aux_loss_coef == 0.02
+
+    def test_no_aux_loss_for_mlp_model(self):
+        """Pure MLP model: model outputs no aux_loss → no MoE metrics."""
+        from astrai.config.model_config import AutoRegressiveLMConfig
+
+        mlp_config = AutoRegressiveLMConfig(**{**TINY_CONFIG, "ffn_type": "mlp"})
+        mlp_model = AutoRegressiveLM(mlp_config).to(self.device)
+        mlp_model.train()
+
+        strategy = SEQStrategy(
+            mlp_model,
+            self.device,
+            moe_aux_loss_coef=0.01,
+        )
+        output = strategy.compute_loss_output(self._make_batch())
+        metrics = output["metrics"]
+
+        assert "moe_aux_loss" not in metrics
+        assert "moe_aux_loss_weighted" not in metrics
+        assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
+        assert strategy._moe_metrics == {}
+
+
+# ── SFTStrategy integration tests ────────────────────────────────────
+
+
+class TestSFTStrategyMoE:
+    """End‑to‑end tests for SFTStrategy with MoE aux loss."""
+
+    @pytest.fixture(autouse=True)
+    def setup(self, device):
+        self.device = device
+        self.config = _make_tiny_moe_config()
+        self.model = _make_model(self.config).to(device)
+        self.model.train()
+
+    def _make_batch(self, batch_size=2, seq_len=8):
+        vocab = self.config.vocab_size
+        input_ids = torch.randint(0, vocab, (batch_size, seq_len))
+        target_ids = torch.randint(0, vocab, (batch_size, seq_len))
+        position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1)
+        loss_mask = torch.ones(batch_size, seq_len, dtype=torch.bool)
+        return {
+            "input_ids": input_ids,
+            "target_ids": target_ids,
+            "position_ids": position_ids,
+            "loss_mask": loss_mask,
+        }
+
+    def test_compute_loss_output_with_aux_loss(self):
+        """SFTStrategy produces MoE metrics when coef > 0."""
+        strategy = SFTStrategy(
+            self.model,
+            self.device,
+            moe_aux_loss_coef=0.01,
+        )
+        output = strategy.compute_loss_output(self._make_batch())
+
+        metrics = output["metrics"]
+        assert "moe_aux_loss" in metrics
+        assert "moe_aux_loss_weighted" in metrics
+        assert metrics["loss"] > metrics["task_loss"] + 1e-12
+
+        moe_metrics = strategy._moe_metrics
+        assert "router_entropy" in moe_metrics
+        assert "dead_expert_fraction" in moe_metrics
+
+    def test_sft_zero_coef_zeroes_weighted_aux(self):
+        """SFTStrategy with zero coef: weighted aux is zero, loss == task_loss."""
+        strategy = SFTStrategy(
+            self.model,
+            self.device,
+            moe_aux_loss_coef=0.0,
+        )
+        output = strategy.compute_loss_output(self._make_batch())
+        metrics = output["metrics"]
+
+        assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
+        assert metrics.get("moe_aux_loss_weighted") == pytest.approx(0.0, abs=1e-6)
+        # Diagnostics still collected
+        assert strategy._moe_metrics
+        assert "router_entropy" in strategy._moe_metrics