From 9b7e6c205fcfebbfc7b107406c29fb36356f9558 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 5 Aug 2026 18:12:28 +0800 Subject: [PATCH] feat: add moe auxloss and metrics --- astrai/config/model_config.py | 1 + astrai/model/components/mlp.py | 13 +- astrai/model/transformer.py | 7 +- astrai/trainer/metric_util.py | 20 ++ astrai/trainer/strategy.py | 114 +++++++++- astrai/trainer/train_callback.py | 10 + scripts/tools/train.py | 8 + tests/module/test_forward_configs.py | 97 ++++++++ tests/module/test_model_forward.py | 316 +++++++++++++++++++++++++++ 9 files changed, 583 insertions(+), 3 deletions(-) create mode 100644 tests/module/test_model_forward.py diff --git a/astrai/config/model_config.py b/astrai/config/model_config.py index b982769..8ad7c83 100644 --- a/astrai/config/model_config.py +++ b/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: diff --git a/astrai/model/components/mlp.py b/astrai/model/components/mlp.py index d0f1401..e76c5bc 100644 --- a/astrai/model/components/mlp.py +++ b/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 diff --git a/astrai/model/transformer.py b/astrai/model/transformer.py index ce069db..7d0df00 100644 --- a/astrai/model/transformer.py +++ b/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) diff --git a/astrai/trainer/metric_util.py b/astrai/trainer/metric_util.py index 8b89bbe..2ec546d 100644 --- a/astrai/trainer/metric_util.py +++ b/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") diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py index 61573f4..4e85443 100644 --- a/astrai/trainer/strategy.py +++ b/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__( diff --git a/astrai/trainer/train_callback.py b/astrai/trainer/train_callback.py index bb29bdd..7cf6842 100644 --- a/astrai/trainer/train_callback.py +++ b/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): diff --git a/scripts/tools/train.py b/scripts/tools/train.py index 5c1dc54..452ca92 100644 --- a/scripts/tools/train.py +++ b/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) diff --git a/tests/module/test_forward_configs.py b/tests/module/test_forward_configs.py index 3f16a9e..4bcb9fd 100644 --- a/tests/module/test_forward_configs.py +++ b/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 diff --git a/tests/module/test_model_forward.py b/tests/module/test_model_forward.py new file mode 100644 index 0000000..71c52e7 --- /dev/null +++ b/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