Exact Quantile Balancing and Load-Error Injection for Mixture-of-Experts
Pit Neitemeier, Jiaze Li, Alessio Serra, Philipp Scholl, Sohir Maskey
cs.LG, cs.CL
2026-09-23
两趟 256 桶全归约恢复 MoE 精确全局分位，局部负载误差直注路由梯度；全局失衡降 19.4%，局部比 GShard 损失低 25%。
MoE 训练里的「负载均衡」是两个问题。全局均衡看一个 optimizer step 内的全部 token：路由长期偏爱少数专家，其余专家学不到东西，参数白养。局部均衡看 expert-parallel 的单个 microbatch：各卡分到的专家负载一歪，通信与计算就要互相等。两个尺度不重合，microbatch 之间的失衡可以在全局上互相抵消，全局均衡了局部照样可以歪。
bias 路线的 Quantile Balancing（QB，出自苏剑林的博客）免辅助损失：给每个专家一个只影响选择的偏置，取值等于路由边距的 K/E 分位数再中心化。麻烦在分布式：全局 batch 切散在各 data-parallel rank 上，谁也看不到完整分布。OpenAthena 的 MARIN 实现各 rank 算本地分位再平均，可分位数与平均不可交换，结果依赖 batch 怎么切；Kimi K3 聚合全局直方图，与切分无关，精度却被桶宽卡住。偏置又是全 batch 共享的一个数，管不了各 microbatch 各自的歪，局部只能靠 GShard 辅助损失这类梯度信号，而它把负载误差经由归一化路由概率传播，专家之间互相耦合。Aleph Alpha 这篇把两个尺度拆开，各给一个干净的解法。
EQB 回答「精确分位数怎么分布式地算」。BF16 的位模式保序，基数选择于是拆成两趟：第一趟各 rank 对边距的高字节做 256 桶计数，all-reduce 后按累积和定位第 r 小落在哪个桶；第二趟只对高字节命中该桶的边距数低字节，再一次 all-reduce，两个字节拼出精确的全局分位。通信量是每层 2·E·256 个 int32 计数，与 token 数无关：E=256 时每层 0.5 MiB，48 个 MoE 层一个 step 合计 24 MiB，比全量 BF16 直方图小 128 倍，还删掉了桶宽这个超参。
LEI 回答「局部误差怎么传给路由」。top-K 是离散的，硬负载目标本身零梯度，GShard 损失与 LEI 都是同一目标的 straight-through 估计，差别只在替身的雅可比：GShard 用归一化概率，把所有专家的残差搅在一起，幅度还被 token 当前的分数和缩放；LEI 改用 batch 均值的原始 sigmoid 分数，雅可比是常数对角阵，修正于是退化成一行：对局部 batch 里每个 token，给专家 e 的分数梯度加上 η·ρe，ρe 是该专家负载偏离均匀值的倍数。梯度下降后，超载专家的分数被压低，欠载的被抬高，前向路由与混合权重都不动。早期实验里两种方法的无界梯度都会把 sliding-window attention 的 LSE 推飞，最终版用 tanh 把残差有界化，保方向、保零中心，这就是表格里的 norm. 变体。
试验台是 7.5B 参数、0.5B 激活的 decoder MoE：256 个路由专家取 6，加 1 个共享专家，50 层里 48 层是 MoE；六个 100B token 消融，另有两个 500B token 对照。失衡指标 MaxVio 取最歪一层里最忙专家超出均匀负载的偏差，局部 MaxVio 5.96 相当于最忙专家扛着约 7 倍于均匀份额的 token。
| 方法 | 全局 MaxVio | 局部 MaxVio | 六基准 BPB 均值 |
| Rank-avg. QB | 0.92 | 5.96 | 0.6782 |
| EQB | 0.74 | 5.38 | 0.6570 |
| EQB + norm. GShard | 0.74 | 4.30 | 0.6758 |
| EQB + norm. LEI | 0.60 | 3.52 | 0.6720 |
BPB 为每字节比特数，越低越好。EQB 换掉 rank 平均 QB：全局失衡降 19.4%，局部降 9.8%，六个基准 BPB 全部变好（MMLU 0.9007→0.8650，ARC 0.7508→0.6876）。按论文正文的对照（对 GShard 损失 0.71/4.71），norm. LEI 全局低 15%、局部低 25%，BPB 均值 0.6720 对 0.6853。500B token 下加 norm. LEI 把局部 MaxVio 从 5.14 压到 4.15，平均准确率从 48.25% 升到 48.56%，代价是全局从 0.44 涨到 0.63。梯度权重 η 是个明确的旋钮：5×10⁻⁶ 时平均准确率最高（50.53%）但全局失衡最差，8×10⁻⁵ 时局部最平（3.54）、准确率回落到 48.00%。
这篇是 aux-loss-free 均衡这条线（DeepSeek-V3 偏置控制器、QB、K3 直方图）的工程收口：把 K3 的近似直方图换成精确值，通信量不涨，顺手删掉桶宽超参。LEI 是 GShard 辅助损失的可平替升级，不改架构、不动前向，只换路由梯度的来源，在训 MoE 的团队可以低成本试。双尺度拆开治理的框架也值得抄：全局用偏置、局部用梯度，互不抢活。
论文自认三条：单一规模、单一种子，没有方差估计；GShard 与 LEI 的系数分别调过，两者梯度量级本来就不同，对比公平性打折；Local MaxVio 只是 dispatch 成本的代理，没有实测 expert-parallel 吞吐或 step 时间。500B token 下 LEI 恶化全局均衡（0.44→0.63），两个尺度确实在互相拉扯。读下来再补一条：100B 消融里 BPB 差距不大（0.6570 对 0.6720），单种子下质量差异不宜过度解读，这篇更硬的证据在均衡指标与通信量上。