分布式训练为什么一加 MoE 负载均衡损失就开始专家更均匀却效果更差：从 Aux Loss 到 Router z-Loss 的工程实战

很多团队把MoE训练不稳归因于专家不均衡于是第一反应是把load balance loss拉高。⚠️ 监控面板很快会变好看专家利用率更均匀热点卡降温token drop也会先降一点。 可真正到上线前复盘时hard set、长链路生成和工具调用样本却先掉点。 这类退化最误导人的地方在于它不表现为训练崩掉而是把路由器从“会挑专家”慢慢推成“只会平均分流”。图 1MoE 监控最先变好看的通常是专家利用率最晚暴露的问题却是专项能力退化专家更平均不等于路由更正确很多实现把Aux Loss当成均衡开关系数一提就想解决热点专家。 但top-k router真正在做的是 token 到专家的判别如果辅助损失过早压过主任务梯度router 会优先追求“每个专家都吃到差不多的 token”而不是“让最合适的专家吃到最该吃的 token”。 结果就是直方图更漂亮专家specialization却被洗平尤其在代码、工具、长样本这类少数模式上最明显。更隐蔽的问题在Router z-Loss和capacity factor的配合。 没有z-loss时router logit 容易在中后段飙大少数专家突然吸走整批 token可只补z-loss不调aux又会把分布压平却保不住判别边界。️ 真正拖垮效果的不是“专家没有平均”而是负载均衡目标和主任务目标在同一批 step 上抢方向最后把稀有模式的专家记忆磨掉。图 2同样是 MoE 均衡Aux Loss 介入时机不同路由器会像两套系统一组 64 卡回放里救效果的不是更强均衡而是把 Aux Loss 推迟生效这次回放使用64卡H800、8x7B MoE结构、全局 batch4096、序列长度8192的指令微调任务。 基线组只开 router 主损失方案二把aux coef固定在0.01方案三把aux线性 warmup 到0.003并加入轻量router z-loss。结果很直接最均匀的一组不是效果最好的一组。方案训练吞吐Token Drop专家熵验证 lossHard set 通过率典型现象仅主任务 router1.00x4.8%1.721.9879.9%热点专家明显aux 0.01固定0.97x2.1%2.082.0575.6%专家均匀但掉点aux warmup z-loss0.99x2.6%1.941.9680.8%均衡与效果更稳真正该盯的不是平均利用率而是 hard slice 的恢复速度。 固定大系数方案在前1.5 kstep 看起来最健康因为专家热度最平可到中后段代码补全、长答案格式和多工具样本先后掉线。✅ 把aux推迟到 router 已学出初步分工后再生效外加一个很轻的z-loss把 logit 尺度钉住效果会比单纯追“更平均”稳得多。defrouter_regularizer(router_logits,load_balance,step):aux_target0.003warmup_steps1800aux_coefaux_target*min(step/warmup_steps,1.0)z_coef1e-4z_losstorch.logsumexp(router_logits,dim-1).pow(2).mean()returnaux_coef*load_balancez_coef*z_loss图 3稳住 MoE 路由的关键不是压到最平均而是让均衡目标晚一点、轻一点介入真正该控的是负载均衡的节奏而不是单看专家熵生产里更稳的顺序通常是先看token drop和热点专家是否伤到吞吐再看aux何时介入最后才看是不是要改top-k或capacity factor。 如果一开始就把aux拉满router 还没学会任务分工就会被迫平均派单如果只看专家熵不看 layer 级别的 hard slice 和长尾样本通过率就很容易把“专家变均匀”误判成“训练变健康”。 尤其在多任务混训里均衡损失过强最先抹掉的往往正是稀有但高价值的专家分工。笔者认为未来3到6个月MoE 训练栈会从“比谁更快做均衡”转向“比谁更会做自适应路由门禁”。 真正能留下来的方案会把aux调度、z-loss、token drop、层间专家熵和专项评测绑在一起而不是只展示一张更平均的利用率曲线。 你们现在更担心的是热点专家拖慢吞吐还是专家看起来更均匀后长尾能力已经在评测集里悄悄下滑如果这篇文章对你有帮助欢迎点赞、收藏和关注。图 4MoE 负载均衡真正要服从的不是图表好看而是主任务学习节奏