refactor: emit strategy metrics as floats

- Converts detached strategy metrics before returning loss output
- Removes redundant item conversion from the trainer loop
- Updates the documented contract and regression tests
This commit is contained in:
2026-08-02 06:38:28 +08:00
parent 1c7369f293
commit 020e2eff4e
4 changed files with 10 additions and 14 deletions
+5 -6
View File
@@ -1,5 +1,6 @@
from types import SimpleNamespace
import pytest
import torch
from astrai.model.transformer import AutoRegressiveLM
@@ -33,16 +34,14 @@ def test_seq_strategy_combines_and_reports_moe_aux_loss(device):
"moe_aux_loss",
"moe_aux_loss_weighted",
}
torch.testing.assert_close(
output["loss"],
assert output["loss"].item() == pytest.approx(
output["metrics"]["task_loss"] + output["metrics"]["moe_aux_loss_weighted"],
)
torch.testing.assert_close(
output["metrics"]["moe_aux_loss_weighted"],
assert output["metrics"]["moe_aux_loss_weighted"] == pytest.approx(
0.25 * output["metrics"]["moe_aux_loss"],
)
assert output["loss"].requires_grad
assert all(not metric.requires_grad for metric in output["metrics"].values())
assert all(isinstance(metric, float) for metric in output["metrics"].values())
def test_metric_callback_includes_dynamic_strategy_metrics(tmp_path):
@@ -103,4 +102,4 @@ def test_legacy_strategy_tensor_loss_is_normalized():
output = strategy({})
assert output["loss"].item() == 2.0
assert output["metrics"]["loss"].item() == 2.0
assert output["metrics"]["loss"] == 2.0