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

tests/trainer/test_loss_output.py

@@ -0,0 +1,106 @@
+from types import SimpleNamespace
+
+import torch
+
+from astrai.model.transformer import AutoRegressiveLM
+from astrai.trainer.strategy import BaseStrategy, SEQStrategy
+from astrai.trainer.train_callback import MetricCallback
+from tests.helpers import make_tiny_config
+
+
+def test_seq_strategy_combines_and_reports_moe_aux_loss(device):
+    config = make_tiny_config(
+        ffn_type="moe",
+        n_routed_experts=4,
+        n_shared_experts=1,
+        n_activated_experts=2,
+        topk_method="greedy",
+    )
+    model = AutoRegressiveLM(config).to(device=device)
+    strategy = SEQStrategy(model, device, moe_aux_loss_coef=0.25)
+    batch = {
+        "input_ids": torch.randint(0, config.vocab_size, (2, 8), device=device),
+        "target_ids": torch.randint(0, config.vocab_size, (2, 8), device=device),
+    }
+
+    output = strategy(batch)
+    legacy_loss = strategy.compute_loss(batch)
+
+    assert isinstance(legacy_loss, torch.Tensor)
+    assert set(output["metrics"]) == {
+        "loss",
+        "task_loss",
+        "moe_aux_loss",
+        "moe_aux_loss_weighted",
+    }
+    torch.testing.assert_close(
+        output["loss"],
+        output["metrics"]["task_loss"] + output["metrics"]["moe_aux_loss_weighted"],
+    )
+    torch.testing.assert_close(
+        output["metrics"]["moe_aux_loss_weighted"],
+        0.25 * output["metrics"]["moe_aux_loss"],
+    )
+    assert output["loss"].requires_grad
+    assert all(not metric.requires_grad for metric in output["metrics"].values())
+
+
+def test_metric_callback_includes_dynamic_strategy_metrics(tmp_path):
+    callback = MetricCallback(
+        ckpt_dir=tmp_path,
+        save_interval=1,
+        metrics=["loss", "lr"],
+    )
+    context = SimpleNamespace(
+        metrics={"task_loss": 2.0, "moe_aux_loss": 1.0},
+        loss=2.01,
+        optimizer=SimpleNamespace(param_groups=[{"lr": 1e-3}]),
+        val_loss=None,
+        grad_norm=None,
+        grad_snr_tracker=None,
+        world_size=1,
+    )
+
+    metrics = callback._metrics(context, callback.metrics)
+
+    assert metrics == {
+        "loss": 2.01,
+        "lr": 1e-3,
+        "task_loss": 2.0,
+        "moe_aux_loss": 1.0,
+    }
+
+
+def test_metric_callback_only_computes_requested_metrics(tmp_path):
+    def fail_metric(context):
+        _ = context
+        raise AssertionError("unrequested metric was computed")
+
+    callback = MetricCallback(
+        ckpt_dir=tmp_path,
+        save_interval=1,
+        metrics=["loss"],
+    )
+    callback._metric_funcs["grad_snr"] = fail_metric
+    context = SimpleNamespace(
+        metrics={},
+        loss=2.0,
+        world_size=1,
+    )
+
+    metrics = callback._metrics(context, callback.metrics)
+
+    assert metrics == {"loss": 2.0}
+
+
+def test_legacy_strategy_tensor_loss_is_normalized():
+    class LegacyStrategy(BaseStrategy):
+        def compute_loss(self, batch):
+            return torch.tensor(2.0, requires_grad=True)
+
+    strategy = LegacyStrategy(torch.nn.Linear(1, 1), "cpu")
+
+    output = strategy({})
+
+    assert output["loss"].item() == 2.0
+    assert output["metrics"]["loss"].item() == 2.0

tests/trainer/test_online_strategy.py

@@ -159,21 +159,21 @@ def test_call_without_runner_falls_back_to_compute_loss_grpo(device):
         "masks": torch.ones(2, 4, 6, device=device),
         "rewards": torch.randn(2, 4, device=device),
     }
-    loss = strat(batch)
+    loss = strat(batch)["loss"]
     assert torch.isfinite(loss).item()
 
 
 def test_call_with_runner_returns_finite_loss_grpo(device):
     strat = _make_grpo(device)
     strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
-    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
+    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})["loss"]
     assert torch.isfinite(loss).item()
 
 
 def test_call_with_runner_returns_finite_loss_dpo(device):
     strat = _make_dpo(device)
     strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
-    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
+    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})["loss"]
     assert torch.isfinite(loss).item()
 
 
@@ -268,7 +268,7 @@ def test_step_called_when_sync_gradients_true(device):
 def test_loss_is_differentiable_dpo(device):
     strat = _make_dpo(device)
     strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
-    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
+    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})["loss"]
     loss.backward()
     has_grad = any(
         p.grad is not None and p.grad.abs().sum() > 0 for p in strat.model.parameters()
@@ -279,7 +279,7 @@ def test_loss_is_differentiable_dpo(device):
 def test_ref_model_not_updated_by_backward_dpo(device):
     strat = _make_dpo(device)
     strat.set_rollout_runner(_RecordingRunner(_make_rollout_result(device=device)))
-    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})
+    loss = strat({"input_ids": torch.randint(3, 200, (2, 4), device=device)})["loss"]
     loss.backward()
     for p in strat.ref_model.parameters():
         assert p.grad is None