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
@@ -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)