Exact quantiles plus local load-error injection: the busiest expert drops from 6.4× to 4.5× its even share

Load balancing is split into two scales: global balance across a full optimizer step so every expert receives enough gradient signal, and local balance within each microbatch so no device becomes a bottleneck under expert parallelism. Earlier approaches either add an auxiliary loss so the router learns balance through gradients, or adjust each expert's routing bias outside gradient descent by a fixed step; the latter does not explicitly account for displaced assignments and approaches balance only gradually. On a 7.5B-parameter proxy with 0.5B active parameters per token trained for 100B tokens, EQB lowered global maximal violation from 0.92 to 0.74 and gave lower mean BPB than rank-averaged QB across six benchmarks.

EQB reads quantile balancing as threshold routing: with each token's bias held fixed, every expert's load can be adjusted independently, so biases come directly from quantiles and need no update rate. Averaging quantiles computed per rank is cheaper but does not generally recover the global quantile and depends on how tokens are partitioned; histogram accumulation has token-independent communication but its quantile accuracy is limited by bin width. EQB encodes each BF16 token-adjusted logit as an order-preserving 16-bit key and uses two-pass radix selection, locating the high byte then resolving the low byte, for 0.5 MiB of reduction payload per layer and 24 MiB across 48 MoE layers per optimizer step, independent of token count.

LEI skips the auxiliary loss and adds each expert's load error on the local microbatch directly to the gradients of its sigmoid score, lowering overloaded experts' scores and raising underloaded ones'. The GShard-style loss uses scores normalized across experts as a differentiable surrogate for discrete loads, so changing one score moves every normalized score for that token; LEI uses unnormalized sigmoid scores, so each score receives only its own expert's load error. Unbounded corrections made attention log-sum-exp grow sharply, and rescaling the whole vector by one shared factor kept it stable; bounded LEI lowered local maximal violation from 5.38 to 3.52, the busiest expert from 6.4× to 4.5× its even share, with mean BPB rising from 0.6570 to 0.6720.

The scope is the load-balancing recipe used in Kolibri pre-training and six comparisons on a 7.5B-parameter proxy with 0.5B active parameters per token, 48 MoE layers, 256 routed experts and 6 active experts per layer, trained for 100B tokens. Evaluation covers MMLU, ARC, HellaSwag, HumanEval, MBPP and GSM8K, reporting mean BPB across the six; attention LSE is a stability check, not a score. Whether to adopt bounded LEI depends on how much local imbalance limits throughput in a given expert-parallel setup.

EQB removes the distributed quantile approximation, but quantization and the limitations of a single QB update remain; biases are computed from the current batch and applied to the next, whose logits differ because data and router parameters changed, so next-batch loads are balanced only approximately. Exactness costs memory: the full-resolution histogram takes 64 MiB per layer with 256 experts, independent of batch size, versus under 1 MiB for Kimi's 1,000-bin histogram. Bounded LEI improves local balance while raising mean BPB by 0.015, and EQB alone has the lowest mean BPB.