refactor: deduplicate low-risk code paths

2026-08-19 16:17:40 +08:00
parent f252af495c
commit 398e8a3ea3
12 changed files with 64 additions and 114 deletions
@@ -94,21 +94,5 @@ def ctx_get_grad_snr(ctx):
    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")
+def ctx_get_moe_metric(ctx, key):
+    return ctx.strategy._moe_metrics.get(key)