Least-Loaded Expert Parallelism: Load Balancing An Imbalanced Mixture-of-Experts

TL;DR

Least-Loaded Expert Parallelism (LLEP) tackles the inefficiencies of Expert Parallelism under imbalanced routing in Mixture-of-Experts models by dynamically redistributing excess workload and expert weights to underutilized devices while preserving exact MoE computation. It introduces a least-loaded assignment (LLA) mechanism with constraints that govern when and how to spill work, along with backward-pass compatibility. Theoretical latency and memory analyses, plus extensive empirical evaluations, demonstrate up to 5–6x speedups on MoE layers and up to 4–5x peak-memory reductions, with end-to-end throughput gains up to 2.2x on large models like gpt-oss-120b. The work provides a principled, hardware-aware framework for tuning EP configurations to maximize throughput in real-world deployments.

Abstract

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.

Figures (9)

• Figure 1: LLEP vs. standard expert parallelism (EP). (a) & (b) show the speedup and peak memory usage per GPU of an MoE layer (128 experts, 4 active experts, hidden size of 2048) under perfectly balanced case and various imbalance scenarios: 30%, 50%, 80%, or 95% of tokens concentrated into 16, 4, 1 imbalanced experts. LLEP is faster than EP by 5$\times$ under extreme imbalance scenarios, while keeping memory usage stable and avoiding out-of-memory risk. (c) Realistic full-model throughput: up to 2.2$\times$ for gpt-oss-20b and 1.9$\times$ for gpt-oss-120b.
• Figure 2: Comparison of standard Expert Parallelism and LLEP.
• Figure 3: Expert routing imbalances across all layers of gpt-oss-20b across batches of a math dataset. (a) E11 has up to 20% load vs. $\sim$3% balanced. (b) GPU 0 has 30-35% vs. $\sim$12.5% balanced. Note that the load numbers do not add up to 100% because values are maximums across all layers.
• Figure 4: Performance comparison of LLEP vs. standard EP across three MoE architectures. Top row: Speedup ($\times$, higher is better) of LLEP over standard EP. Gray bars show the balanced baseline ($\approx$1$\times$). Colored bars indicate imbalance levels (percentage of tokens routed to that many experts). Higher concentration yields greater speedup, up to 6.1$\times$ for GPT-OSS-120B. Bottom row: Peak memory usage per GPU (GB, lower is better). Hatched bars represent standard EP; solid bars represent LLEP. EP memory grows dramatically with imbalance (up to 100GB for Kimi-K2), while LLEP maintains near-constant memory across all scenarios.
• Figure 5: Performance (accuracy on AIME'25) vs. wall-time for EP and LLEP when training gpt-oss-20b (low effort) on full parameters using Zero-3 and CPU offloading for gradients and optimizer states. CPU operations and checkpoint saving introduce non-negotiable overheads. Despite additional overhead, training converges 1.25x faster with LLEP as compared to EP.