PaperlayerView exact revision on arXiv

Source-linked AI summary

Exact Quantile Balancing and Load-Error Injection for Mixture-of-Experts

Pit Neitemeier, Jiaze Li, Alessio Serra, Philipp Scholl, Sohir Maskey

arXiv:2609.28053v1cs.LGcs.CL

TL;DR

MoE training needs global balance across optimizer steps and local balance within expert-parallel microbatches, while existing distributed QB has shard-dependent or approximate global quantiles and shared biases cannot adapt to microbatches. The paper introduces EQB for exact global quantiles and LEI for local load-error correction. EQB improves over rank-averaged QB, and LEI improves balance over GShard at comparable quality.

  • Problem

    Global and local MoE load balance are distinct, while distributed QB may use shard-dependent or approximate quantiles and a shared token-independent bias cannot adapt to local microbatches.

    Source: §1 · p. 1

  • Method

    EQB uses two-pass BF16 radix selection for exact global-batch quantiles, while LEI injects local load errors directly into router-score gradients.

    Source: §3 · pp. 2–3, §4 · p. 3

  • Results

    EQB improves on rank-averaged QB, while LEI improves balance over GShard at comparable quality.

    Source: §6 · p. 4

  • Takeaways & Limitations

    EQB and LEI provide complementary control of global and local MoE load balance.

    Source: §6 · p. 4

    Results use one model scale and seed without variance estimates, and Local MaxVio is a proxy rather than measured EP throughput.

    Source: §6 · p. 4

Abstract

from arXiv · show

Mixture-of-Experts (MoE) training requires global load balance to prevent expert under-utilization and local balance for efficient expert-parallel execution. Existing distributed Quantile Balancing (QB) uses shard-dependent or approximate global quantiles, while token-independent expert biases cannot ensure microbatch-level balance. We introduce Exact Quantile Balancing (EQB), which computes exact global-batch BF16 quantiles with negligible communication, and Load-Error Injection (LEI), which injects local load errors directly into router-score gradients. On 7.5B-parameter MoEs trained for up to 500B tokens, EQB improves global balance and downstream performance over naive QB, while LEI improves local balance and outperforms the GShard loss at comparable quality.

1 Introduction

MoE training needs both optimizer-step global balance and expert-parallel microbatch local balance, which are distinct objectives. The paper introduces EQB for exact global quantiles and LEI for local load-error correction, reporting improvements over prior approaches.

1 Introduction

  • Global balance prevents persistent concentration on a few experts, while local balance reduces microbatch load skew for dispatch and expert-compute efficiency.Imbalances across local microbatches can cancel over an optimizer step, so global balance does not imply local balance.Source: §1 · p. 1
  • EQB recovers the exact global-batch empirical quantile with two-pass BF16 radix selection and two 256-bin all-reduces per layer.Distributed QB alternatives use shard-averaged quantiles that depend on batch partitioning or histogram estimates limited by their resolution.Source: §1 · p. 1
  • LEI injects each expert’s local load error directly into router-score gradients; experiments report better local balance than GShard at comparable quality.The paper reports that LEI can improve local balance and average accuracy together.Source: §1 · p. 1

2 Preliminaries: MoE Routing and Load Balancing

The section defines MoE routing and load measures, then contrasts global quantile-based bias updates with local gradient-based balancing. The two approaches address different load scales and have different distributed-training constraints.

2 Preliminaries: MoE Routing and Load Balancing

  • An MoE router selects K experts from logits shifted by expert bias, but weights selected experts using unbiased sigmoid scores.The router forms z_t = W_r x_t and selects I_t = Top-K(z_t + b).Source: §2 · p. 2
  • Expert load counts assignments, with target K|B|/E; the global batch spans an optimizer step, while a local batch is one rank-local EP microbatch.The hard-load fraction normalizes an expert’s assignment count by K|B_l|.Source: §2 · p. 2
  • QB sets expert bias from routing-margin quantiles to target selection by a fraction K/E, then centers the update; the resulting bias is used in the next step.The centering removes a common offset that does not affect top-K selection.Source: §2 · p. 2
  • Averaging rank-local quantiles generally differs from the global-batch quantile, while histogram aggregation is partition-invariant but approximate at its chosen resolution.These constraints arise because quantiles do not commute with averaging and histogram precision depends on bin width.Source: §2 · p. 2
  • The GShard auxiliary loss uses local hard-load fractions and normalized routing probabilities to encourage reduced load imbalance through router gradients.The hard-load fraction is treated as stop-gradient; gradient-based control is complementary to bias-based methods.Source: §2 · p. 2

3 Exact Quantile Balancing

Exact global QB requires an order statistic across margins distributed over data-parallel ranks, but gathering all margins scales communication with token count. EQB uses two-pass BF16 radix selection to recover that statistic exactly with communication independent of token count.

3 Exact Quantile Balancing

  • EQB adapts radix selection to locate the required global order statistic without gathering all token-level margins or building an exact histogram over all BF16 values.Direct gathering communicates with token count, while an exact BF16 histogram could require up to 2^16 bins per expert.Source: §3 · p. 2
  • Two 256-bin passes over the high and low BF16 bytes recover the exact global-batch empirical quantile while keeping token-level margins local.All-reduced counts identify the target high-byte bin and then the target low-byte rank within it.Source: §3 · p. 3
  • EQB communicates 2E·256 int32 counts per layer, independent of token count, and returns the exact empirical quantile.For E = 256, the passage reports 0.5 MiB per layer and 24 MiB per optimizer step for 48 MoE layers.Source: §3 · p. 3

4 Load-Error Injection

Load-Error Injection (LEI) uses local hard-load error to correct router-score gradients, with a constant, expert-specific surrogate rather than normalized routing probabilities. Its injected residual can be bounded without changing its direction, and forward routing remains unchanged.

Load-Error Injection

  • The local objective compares each expert’s hard-load fraction F_e with the uniform target Q_e = 1/E.The hard-load residual F − Q determines the desired correction, while the surrogate determines how it reaches router scores.Source: §4 · p. 3
  • Unlike GShard’s normalized-probability surrogate, LEI avoids coupling an expert’s correction to other experts’ residuals.GShard’s probability-dependent Jacobian can change without changing hard top-K assignments.Source: §4 · p. 3
  • LEI uses the mean unnormalized router score as its straight-through estimator, giving a constant diagonal Jacobian.Each expert’s load error affects only that expert’s router scores and is applied equally across local-batch tokens.Source: §4 · p. 3
  • A shared scale bounds LEI’s injected residual while preserving its direction and zero-centering.LEI leaves forward routing and mixture weights unchanged.Source: §4 · p. 4

5 Experiments

Experiments compare EQB and LEI variants on a 7.5B-parameter MoE, measuring global and local load imbalance. EQB improves balance and benchmark BPB over rank-averaged QB, while LEI improves local balance relative to GShard at comparable quality in the 100B-token ablation and at 500B tokens with a small increase in Global MaxVio.

Setup

  • The evaluation uses a 7.5B-parameter decoder-only MoE with 0.5B active parameters per token, 256 routed experts, and top-6 routing.Global MaxVio aggregates loads over an optimizer step; Local MaxVio is measured on rank-local expert-parallel microbatches.Source: §5 · p. 4

Results

  • EQB reduces Global MaxVio from 0.92 to 0.74 (19.4%) and Local MaxVio from 5.96 to 5.38 (9.8%) versus rank-averaged QB, while improving BPB on all six reported benchmarks.The comparison replaces rank-averaged quantiles with the exact global-batch statistic.Source: §5 · p. 4
  • At 500B tokens, LEI reduces Local MaxVio from 5.14 to 4.15 as average accuracy rises from 48.25% to 48.56%, while Global MaxVio increases from 0.44 to 0.63.These are the reported comparison values for EQB and normalized LEI.Source: §5 · p. 4
  • At 100B tokens, normalized LEI reaches Global/Local MaxVio of 0.60/3.52 versus 0.71/4.71 for GShard, at comparable model quality.Six ablations train for 100B tokens; EQB and normalized LEI are also compared at 500B tokens.Source: §5 · p. 4

6 Conclusion

EQB and LEI target complementary scales of MoE load balancing, while related work motivates separating global bias control from local score-side correction. The conclusion reports improved balance under the stated conditions, with limitations in model scale, seed coverage, and the use of Local MaxVio as a throughput proxy.

6 Conclusion

  • EQB and LEI address complementary scales: EQB recovers the exact global-batch quantile, while LEI injects local load residuals into router-score gradients.The conclusion reports EQB improves on rank-averaged QB and LEI improves balance over GShard at comparable quality; bounding LEI’s residual avoids attention-logit instability.Source: §6 · p. 4
  • Results use one model scale and seed without variance estimates, and Local MaxVio is a proxy for dispatch cost rather than measured EP throughput.GShard and LEI use separately tuned coefficients because the same coefficient produces very different router-gradient magnitudes.Source: §6 · p. 4

A Related work

  • GShard couples hard expert loads to mean routing probabilities, while Wang et al. [2024] update selection biases from observed loads using a tuned step size.The paper uses these approaches to motivate separating global bias control from local score-side correction.Source: A Related work · p. 6
  • Distributed QB averages rank-local quantiles or estimates the global quantile with a histogram, whereas EQB targets the global-batch statistic exactly.LEI complements token-independent QB bias with a local router-score update.Source: A Related work · p. 6

C Exact Quantile Balancing Algorithm

  • The EQB algorithm is presented as a two-pass procedure that encodes BF16 routing margins into order-preserving high and low bytes.The described coarse pass uses the high byte, and the listed count records entries for a selected high-byte value by low-byte value.Source: C Exact Quantile Balancing Algorithm · p. 7

D Straight-through derivation of the GShard loss and LEI

The derivation frames load-balancing gradients as an observed hard-load residual multiplied by a surrogate Jacobian. GShard and LEI share this hard-load objective but use different surrogate Jacobians, while LEI applies the residual directly to router-score gradients.

D Straight-through derivation of the GShard loss and LEI

  • A stop-gradient construction keeps the observed hard load in the forward pass and substitutes a differentiable surrogate’s gradient in the backward pass.The resulting gradient is the surrogate Jacobian applied to the hard-load residual.Source: D Straight-through derivation of the… · p. 8
  • The GShard surrogate uses batch-mean normalized routing probabilities, coupling each expert’s residual to the other experts through a probability-weighted mean.Its score-dependent Jacobian also scales with the token’s current score sum.Source: D Straight-through derivation of the… · p. 8
  • With an expert bias, the estimator combines loads selected using z + b with probabilities computed from s alone, which can weaken alignment when the bias changes top-K choices.The score-dependent Jacobian can further attenuate the correction.Source: D Straight-through derivation of the… · p. 8
  • LEI uses a batch-mean unnormalized score with a constant diagonal Jacobian, passing the same observed load residual to every token up to coefficient and residual rescaling.GShard and LEI are STEs of the same hard-load objective but differ in their surrogate Jacobians; the expert bias determines the observed load and is held fixed during differentiation.Source: D Straight-through derivation of the… · pp. 8–9
  • A gradient-based auxiliary-loss-free bias controller uses the same identity-like surrogate as LEI but acts on token-independent selection bias, whereas QB uses an exact coordinate-wise quantile minimizer.The controller’s residual is F − Q.Source: D Straight-through derivation of the… · p. 9

E.1 Model, training, and evaluation

The experiments use a 7.5B-parameter decoder-only MoE and include six 100B-token ablations plus longer 500B-token runs. Evaluations cover eight benchmarks, and balance metrics are computed across layers and training steps.

E.1 Model, training, and evaluation

  • The decoder-only model has 7.5B parameters, 0.5B active per token, 50 Transformer layers, and 48 MoE layers with 256 routed experts and top-6 routing.It uses a shared expert, sigmoid routing, a 128k vocabulary, 16,384-token sequences, and sliding-window to full attention in a 4:1 ratio.Source: §E.1 · p. 9
  • Six ablations train for 100B tokens over 10,342 steps, while EQB and normalized LEI are additionally trained for 500B tokens.The ablations use 64 data-parallel ranks, a global batch of 576 sequences, and a warmup-stable learning-rate schedule.Source: §E.1 · p. 9
  • Evaluation covers MMLU, ARC, HellaSwag, TriviaQA, GSM8K, HumanEval, MBPP, and MATH using the stated framework and Minerva protocol.Global MaxVio aggregates loads over an optimizer step; Local MaxVio averages per-microbatch and rank maxima, with each metric taking the maximum across layers and averaging the final 1,000 steps.Source: §E.1 · p. 9
  • The experiments also track maximum pre-softmax attention-logit LSE across SWA layers, tokens, and heads as an indicator of attention-logit instability.Source: §E.1 · p. 9

E.2 100B-token ablation

The section points readers to training trajectories, endpoint results, and per-benchmark evaluation.

E.2 100B-token ablation

  • Figure 1 shows training trajectories corresponding to Table 1 endpoint results, while Table 3 reports the complete per-benchmark evaluation.Source: §E.2 · p. 9

E.3 500B-token results

The 500B-token evaluation includes a normalized-LEI gradient-weight sweep and corresponding all-benchmark BPB results.

E.3 500B-token results

  • At 100B tokens, LEI reduces both Global and Local MaxVio more strongly than the GShard loss; normalization keeps LSE stable while preserving most balance gains.Source: Figure 1 · p. 10
  • Table 3 presents a six-run 100B-token ablation across all benchmarks, with results reported as BPB, where lower is better.Source: Table 3 · p. 10
  • Table 4 reports a normalized-LEI gradient-weight sweep at 500B tokens, with balance metrics averaged over the final 1,000 steps.Its Avg. is the unweighted mean of reported accuracy metrics.Source: §E.3 · p. 9, Table 4 · p. 10
  • Table 5 covers all benchmarks in the 500B-token setting, reporting BPB, where lower is better.Source: §E.3 · p. 9, Table 5 · p. 11
Loading 2609.28053v1…