Source-linked AI summary
Least-Loaded Expert Parallelism: Load Balancing An Imbalanced Mixture-of-Experts
Xuan-Phi Nguyen, Shrey Pandit, Austin Xu, Caiming Xiong, Shafiq Joty
TL;DR
Standard EP assumes balanced routing, but well-trained MoE models can route tokens unevenly toward specialized experts, causing overloaded devices during post-training and inference. LLEP dynamically reroutes excess tokens and expert weights to least-loaded devices while preserving exact MoE computation. Across model scales, it improves speed and peak memory, including up to 90% higher end-to-end throughput for gpt-oss-120b.
Problem
Standard EP assumes balanced GPU loads, although MoE routing can be persistently imbalanced and explicit load balancing is discouraged during post-training or inference.
Method
LLEP dynamically transfers excess tokens and corresponding expert weights from overloaded devices to underloaded devices while preserving exact MoE computation.
Results
Up to 90% higher end-to-end throughput is reported for gpt-oss-120b, alongside up to 5-6× MoE-layer speedups and 5× lower peak memory consumption.
Takeaways & Limitations
LLEP distributes workload and memory across devices, enabling efficient execution of imbalanced MoE workloads while maintaining the model’s routing behavior.
Takeaways & Limitations
LLEP uses thresholds for GPU capacity, minimum efficient GEMM load, and imbalance detection, and switches back to standard EP when routing is balanced enough.
Abstract
from arXiv · showhide
Mixture-of-Experts (MoE) models are typically pre-trained with explicit load-balancing constraints to ensure statistically balanced expert routing. Despite this, we observe that even well-trained MoE models exhibit significantly imbalanced routing. This behavior is arguably natural-and even desirable - as imbalanced routing allows models to concentrate domain-specific knowledge within a subset of experts. Expert parallelism (EP) is designed to scale MoE models by distributing experts across multiple devices, but with a less-discussed assumption of balanced routing. Under extreme imbalance, EP can funnel a disproportionate number of tokens to a small number of experts, leading to compute- and memory-bound failures on overloaded devices during post-training or inference, where explicit load balancing is often inapplicable. We propose Least-Loaded Expert Parallelism (LLEP), a novel EP algorithm that dynamically reroutes excess tokens and associated expert parameters from overloaded devices to underutilized ones. This ensures that all devices complete their workloads within the minimum collective latency while respecting memory constraints. Across different model scales, LLEP achieves up to 5x speedup and 4x reduction in peak memory usage compared to standard EP. This enables faster and higher-throughput post-training and inference, with ~1.9x faster for gpt-oss-120b. We support our method with extensive theoretical analysis and comprehensive empirical evaluations, including ablation studies. These results illuminate key trade-offs and enable a principled framework for hardware-specific hyper-parameter tuning to achieve optimal performance.
1 INTRODUCTION
MoE routing can remain substantially imbalanced after training because task-relevant experts specialize, but standard expert parallelism assumes balanced GPU loads. LLEP dynamically redistributes excess tokens and expert weights while preserving exact MoE computation, improving speed and memory use.
- Imbalanced routing can be natural and desirable because experts may specialize in domains or tasks while others generalize.
- Standard EP assumes approximately balanced GPU loads, so extreme routing imbalance can concentrate tokens on a few devices.
- LLEP dynamically transfers excess tokens and corresponding expert weights from overloaded devices to underloaded devices.
- LLEP triggers token transfer only when communication costs less than processing the tokens locally and supports backward-pass execution.
- 1.9× speedup is achieved for gpt-oss-120b, while LLEP reaches up to 4.6× speedup under extreme imbalance and up to 4× lower peak-memory growth than standard EP.LLEP maintains similar throughput to standard EP under balanced routing; the cited results also report up to 1.4× speedup for gpt-oss-20b.
2 BACKGROUND
MoE layers route tokens sparsely through selected feed-forward experts, while expert parallelism distributes those experts across GPUs. Dispatch and combine use sorting, re-indexing, and All-to-All exchanges to process and return routed tokens.
- 2.1 MIXTURE-OF-EXPERTS: An MoE router selects the top-K experts for each token from a set of feed-forward experts with distributed weight matrices.
- 2.1 MIXTURE-OF-EXPERTS: The router computes affinity scores from the token representation and router weights, then retains only selected experts’ scores for the output.
- 2.2 EXPERT PARALLELISM: EP distributes experts across GPUs, with each device hosting only a local subset of experts.
- 2.2 EXPERT PARALLELISM: The dispatch-combine procedure sorts and re-indexes routed tokens, applies local expert computations, then reverses the exchange and ordering.
- 2.2 EXPERT PARALLELISM: During dispatch, devices exchange tokens with GPUs hosting their selected experts through an All-to-All communication operation.
3 ANALYSIS
Well-trained MoE models can exhibit domain-driven routing imbalance, exposing standard expert parallelism to severe latency and memory costs. The analysis motivates distributing expert work while accounting for GEMM efficiency and assignment overhead.
- 3.2 DISTRIBUTED LATENCY AND MEMORY ANALYSIS: Least-loaded assignment partitions expert token loads across GPUs and constructs a corresponding expert-weight transfer plan.The algorithm processes experts from largest to smallest load, using native capacity first and spilling excess work to other devices.
- 3.2 DISTRIBUTED LATENCY AND MEMORY ANALYSIS: GEMM efficiency generally improves with larger token batches and model dimensions, making many small expert GEMMs slower than fewer large GEMMs at equal FLOPs.Kernel-launch overhead and hardware-specific implementation choices also affect the latency trade-off.
- 3.2 DISTRIBUTED LATENCY AND MEMORY ANALYSIS: Standard EP assigns each expert’s globally routed tokens to its resident device, so one expert can receive nearly the entire global batch.This concentration produces spiking latency and memory usage on the overloaded device.
- 3.2 DISTRIBUTED LATENCY AND MEMORY ANALYSIS: 4.6x slower execution and up to 4x higher peak memory occur when 95% of tokens route to one expert under standard EP.The overloaded device may run out of memory while other devices remain underutilized.
4 LEAST-LOADED EXPERT PARALLELISM (LLEP)
LLEP detects routing imbalance and, when it exceeds a threshold, assigns expert-token portions across devices while transferring weights as needed. Its capacity, minimum-GEMM-size, and switching thresholds must be tuned to the model and hardware.
- 4 LEAST-LOADED EXPERT PARALLELISM (LLEP): LLEP switches from standard EP to least-loaded assignment when global routing imbalance exceeds threshold λ.The assignment determines which GPUs compute each expert’s token portions and imports nonresident expert weights when necessary.
- 4 LEAST-LOADED EXPERT PARALLELISM (LLEP): Factor α sets each GPU’s maximum token capacity through mα, and excess local-expert load spills to other GPUs.mα is an overload threshold rather than necessarily a physical memory limit.
- 4 LEAST-LOADED EXPERT PARALLELISM (LLEP): Minimum GEMM size m prevents spilling when the excess workload is too small to process efficiently elsewhere.In that case, the native GPU computes the excess despite exceeding its occupied capacity.
- 4 LEAST-LOADED EXPERT PARALLELISM (LLEP): The imbalance threshold λ avoids LLEP’s assignment overhead when routing is already balanced and standard EP would produce the same plan.The optimal α, m, and λ depend on model dimensions, workload, model size, and physical system configuration.
- 4 LEAST-LOADED EXPERT PARALLELISM (LLEP): The evaluated implementation uses Torch NCCL for All-to-All and peer-to-peer operations, while its LLA algorithm is implemented in pure Python.The authors identify low-level C++/Triton kernels and fused communication as further optimization opportunities.
5 EXPERIMENTS
Experiments evaluate LLEP through controlled imbalance simulations, end-to-end throughput and training comparisons, and hyperparameter ablations. LLEP improves speed and memory behavior under imbalance while preserving standard EP efficiency in balanced settings.
- 5.1 SPEED AND MEMORY PROFILES: Controlled benchmarks evaluate popular MoE layers across gpt-oss-120b, DeepSeek-V3, and Kimi-K2 under balanced and 30%–95% concentrated routing scenarios.Experiments use 8 H200 GPUs, with λ = 1.3, α = 1, and m = 1024 for LLEP.
- 5.1 SPEED AND MEMORY PROFILES: 6.11× speedup is achieved in the most extreme case, with 95% of tokens routed to one expert, while balanced routing retains EP efficiency.LLEP outperforms standard EP across imbalance scenarios and adaptively maintains comparable efficiency when routing is perfectly balanced.
- 5.1 SPEED AND MEMORY PROFILES: 5× memory saving is achieved across imbalance scenarios, as LLEP keeps memory stable while standard EP’s usage grows and can trigger OOM failures.The reduced peak memory can support larger batch sizes without running out of memory.
- 5.2 END-TO-END FULL MODEL SPEED PROFILES IN THE WILD: 1.25× speedup is achieved during full-parameter gpt-oss-20b training with comparable AIME’25 accuracy despite CPU and checkpointing overheads.The comparison uses Zero-3 with CPU offloading for gradients and optimizer states.
- 5.3 ABLATION STUDY: LLEP’s speedup increases with batch size because balanced GPU workloads reduce collective processing time and large-batch communication overshadows weight-transfer overhead.Higher α reduces speedup, while larger hidden sizes improve LLEP’s scaling as GEMM efficiency and GPU saturation increase.
- 5.3 ABLATION STUDY: The adaptive ratio λ favors reverting to standard EP when routing is sufficiently balanced and LLEP’s weight-transfer overhead exceeds the benefit of even computation.At B = 8K, higher λ is beneficial when imbalance is low at 15–20%.
6 CONCLUSION
LLEP dynamically balances imbalanced MoE routing across devices while preserving exact MoE computation. It delivers substantial speedups and peak-memory reductions for MoE layers and improves gpt-oss-120b throughput.
- LLEP dynamically load-balances imbalanced routing while ensuring the exact MoE mathematical computation.
- LLEP achieves up to 5-6× speedups and 5× lower peak memory consumption for MoE layers.
- LLEP improves gpt-oss-120b end-to-end full-model throughput by up to 90%.
A.1 SEPARATE VS FUSED GROUPED-GEMM
Grouped-GEMM execution time increases with the number of experts even when total FLOPs remain fixed. Standard expert parallelism and LLEP mitigate this by distributing expert weights so each rank computes only a few experts.
- Execution time increases with the number of experts despite identical total FLOPs.The benchmark uses 65536 evenly distributed tokens and H = D = 8192.
- Spreading expert weights across EP ranks lets each rank compute only a handful of experts.
A.2 NUMBER OF EXPERTS
LLEP becomes more efficient as the number of experts increases. Its speedup over standard EP is evaluated as a function of expert count with four imbalanced experts.
- LLEP exhibits greater speedups as the number of experts in the MoE layer increases.
- Figure 9 measures LLEP speedup over standard EP as a function of expert count with four imbalanced experts.