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

TL;DR - This paper introduces Exact Quantile Balancing (EQB) and Load-Error Injection (LEI) to improve global and microbatch-level expert balance in distributed Mixture-of-Experts training. Experiments on 7.5B-parameter models indicate better balance and downstream performance with negligible added communication.

• EQB computes exact global-batch BF16 routing quantiles, avoiding shard-dependent or approximate estimates.
• LEI injects local expert-load errors directly into router-score gradients to improve microbatch balance.
• Across training runs of up to 500B tokens, EQB outperforms naive Quantile Balancing in global balance and downstream quality.
• LEI achieves better local balance than the GShard auxiliary loss at comparable model quality.