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
Exact two-pass radix quantiles plus direct load-error gradient injection cut MoE imbalance on a 7.5B model: global down 19.4%, local 25% below the GShard loss.
Load balancing in MoE training is two problems at once. Global balance, measured over all tokens in an optimizer step, keeps routing from concentrating on a few experts and starving the rest. Local balance, inside each expert-parallel microbatch, decides whether devices sit idle waiting for the most-loaded expert. The two do not imply each other: skews across microbatches can cancel over a step, so a globally balanced router can still dispatch badly.
Bias-based Quantile Balancing (QB, from Jianlin Su's blog) avoids auxiliary losses: each expert gets a selection-only bias set to the K/E quantile of routing margins over the global batch. Distributed training breaks the assumption, because no data-parallel rank sees the full margin distribution. The MARIN implementation averages rank-local quantiles, but quantiles do not commute with averaging, so the result depends on how the batch is sharded. Kimi K3 aggregates a global histogram instead: partition-invariant, but its accuracy is capped by bin width. A bias shared across the whole batch also cannot adapt to individual microbatches, which is why local balance still falls back to gradient signals like the GShard auxiliary loss, whose correction flows through normalized routing probabilities and couples all experts. This paper from Aleph Alpha splits the two scales and gives each a clean fix.
EQB answers how to compute the exact quantile in a distributed setting. Because BF16 bit patterns preserve value order, radix selection reduces to two passes over the two bytes: each rank histograms the high byte of its margins into 256 bins, one all-reduce locates the bin holding the r-th smallest value, and a second pass histograms the low byte within that bin. Concatenating the two bytes recovers the exact global quantile. Communication is 2·E·256 int32 counts per layer, independent of token count: 0.5 MiB per layer at E=256, 24 MiB per optimizer step across the 48 MoE layers, 128× less than all-reducing a full BF16 histogram, with no bin-width hyperparameter left to tune.
LEI answers how to send local errors to the router. Top-K is discrete, so the hard-load objective has zero gradient; the GShard loss and LEI are both straight-through estimators of it, differing only in the surrogate Jacobian. GShard's normalized-probability surrogate couples every expert's residual and rescales the correction by each token's current score sum. LEI uses the batch-mean raw sigmoid score, whose Jacobian is constant and diagonal, collapsing the correction to one line: add η·ρe to the score gradient of expert e for every token in the local batch, where ρe is that expert's relative load error. After descent, overloaded experts see their scores pushed down and underloaded ones pulled up; forward routing and mixture weights are untouched. Unbounded versions of both methods blew up sliding-window attention LSE in early runs, so the shipped variant bounds the residual with a tanh rescale that keeps direction and zero-centering.
The testbed is a 7.5B-parameter, 0.5B-active decoder MoE with 256 routed experts (top-6, plus one shared), 48 MoE layers out of 50, six 100B-token ablations and two 500B-token runs. MaxVio reports how far the busiest expert in the worst layer exceeds uniform load; a local MaxVio of 5.96 means roughly 7× the uniform share.
| Method | Global MaxVio | Local MaxVio | Mean BPB over 6 benchmarks |
| 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 is bits per byte, lower is better. EQB over rank-averaged QB: global imbalance down 19.4%, local down 9.8%, BPB better on all six benchmarks (MMLU 0.9007→0.8650, ARC 0.7508→0.6876). Against the GShard loss as reported in the paper's text (0.71/4.71), normalized LEI is 15% lower globally and 25% lower locally, with mean BPB 0.6720 versus 0.6853. At 500B tokens, adding normalized LEI cuts local MaxVio from 5.14 to 4.15 and lifts average accuracy from 48.25% to 48.56%, while global MaxVio rises from 0.44 to 0.63. The gradient weight η is an explicit dial: 5×10⁻⁶ gives the best average accuracy (50.53%) with the worst global balance, while 8×10⁻⁵ gives the flattest local balance (3.54) with accuracy back at 48.00%.
This is the engineering cleanup of the aux-loss-free balancing line (the DeepSeek-V3 bias controller, QB, K3's histogram): swap the approximate histogram for an exact statistic at no extra communication and drop a hyperparameter. LEI is a drop-in replacement for the GShard auxiliary loss: same training graph, different gradient source, better local balance at equal quality. The two-scale split is the part worth copying: bias for the global scale, gradients for the local one, each doing its own job.
The authors flag three: one model scale, one seed, no variance estimates; GShard and LEI coefficients tuned separately, since the two produce very different gradient magnitudes; and local MaxVio is a proxy for dispatch cost, not measured expert-parallel throughput. At 500B tokens LEI worsens global balance (0.44 to 0.63), so the two scales do pull against each other. One more from reading the tables: the BPB gaps in the 100B ablation are small (0.6570 versus 0.6720), and with a single seed the quality differences should not be over-read; the solid evidence here sits in the balance metrics and the communication bound.