feat: add moe auxloss and metrics

2026-08-05 18:12:28 +08:00
parent 602b5ce216
commit 9b7e6c205f
9 changed files with 583 additions and 3 deletions
@@ -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):