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

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.

• The paper introduces a dynamic load-balancing algorithm (LLEP) that reassigns tokens from overloaded to underloaded GPUs in MoE architectures.
• It employs a Least-Loaded Assignment strategy to achieve significant speedups (5-6x) and reduce peak memory usage by up to 5x.
• Empirical results on models like GPT-OSS demonstrate substantial throughput gains while preserving expert specialization.

Least-Loaded Expert Parallelism for Load Balancing in Imbalanced Mixture-of-Experts Architectures

Background and Motivation

Mixture-of-Experts (MoE) architectures have established themselves as a cornerstone for scaling LLMs due to their ability to combine expert specialization with efficient conditional computation. MoE layers distribute token processing across multiple feed-forward experts, with selection mediated by a gating router. Despite the use of auxiliary load-balancing losses and stochastic biases during pre-training, empirical evidence demonstrates that post-training or inference on domain-specific data frequently yields highly imbalanced expert activation. This non-uniform token-to-expert distribution is not only hard to prevent but is also functionally advantageous for domain specialization.

Expert Parallelism (EP) is the prevailing strategy for scaling MoEs across GPUs and nodes, relying on the assumption of uniformly distributed load. However, practical workloads often violate this assumption, resulting in scenarios where a small subset of experts are responsible for the majority of computation, overloading their host devices. Such load concentration leads to compute-bound slowdowns or out-of-memory (OOM) failures during both fine-tuning and inference, especially when explicit rebalancing strategies cannot be deployed.

Least-Loaded Expert Parallelism (LLEP): Methodology

LLEP introduces a dynamic load balancing algorithm for MoE expert parallelism. Unlike earlier methods that rely on parameter-altering regularization or expert replication (which bears significant memory overhead and limited applicability), LLEP remaps excess tokens — and their associated expert weights — from overloaded GPUs to underloaded ones in real time. This is facilitated through a Least-Loaded Assignment (LLA) procedure that greedily assigns computation to available memory/compute resources by considering both per-GPU capacity and communication overhead.

Key constraints and hyperparameters in LLEP include:

• Token Capacity Factor ($\alpha$): Controls the maximum permissible tokens for a GPU before triggering load redistribution.
• Minimum GEMM Tokens ($m$): Ensures that a transferred data chunk is large enough to maintain efficient GEMM operation.
• Imbalance Threshold ($\lambda$): Adapts between standard EP and LLEP depending on the degree of load imbalance, reverting to EP in balanced scenarios.

LLEP is architecture-agnostic, supporting backward gradient propagation and direct drop-in replacement for existing EP infrastructure. Weight transfers accompanying token transfers only occur when the performance gain surpasses the communication cost, supporting both multi-GPU/multi-node scaling and hardware-specific configuration.

Empirical Results and Analysis

Extensive theoretical and practical analyses validate the superior performance of LLEP across multiple model scales, including GPT-OSS-20B/120B, DeepSeek-V3, and Kimi-K2 configurations. Under extreme imbalance scenarios — e.g., 95% of tokens routed to a single expert — LLEP delivers:

• Up to 5-6x speedup in MoE layer computation over standard EP.
• Peak memory consumption reduced by up to 5x, ensuring robust OOM protection and enabling larger batch sizes.
• Full-model throughput gains of up to 1.9x for gpt-oss-120b and 2.2x for gpt-oss-20b in end-to-end inference.

Results are consistent across varying batch sizes, number of experts, and hidden dimensions, with more pronounced improvements at larger model scales and higher degrees of imbalance. Ablation studies show that performance scales linearly with batch size, expert count, and hidden size, highlighting LLEP’s efficacy for modern LLM workloads.

Implications and Future Directions

LLEP shifts the paradigm from insisting on uniformly balanced expert activation during deployment to system-layer load redistribution, maximizing utilization, throughput, and memory safety without altering model behavior. This respects and leverages the specialization learned during pre-training, avoiding the pitfalls of expert collapse but also not interfering with useful specialization. Hardware-conscious hyperparameter tuning is supported, providing a principled framework adaptable to heterogeneous accelerator environments.

The approach opens avenues for scaling future multi-trillion-parameter MoE models, reducing infrastructure cost, and simplifying the deployment pipeline in both research and production. Extensions may include optimizing communication via fused kernels, supporting multi-node routing within and across clusters with varying communication bandwidths, and load-aware global scheduling in distributed training and inference.

Conclusion

Least-Loaded Expert Parallelism delivers a robust solution for load balancing in imbalanced MoE deployments, addressing key practical limitations of existing expert parallelism techniques. LLEP realizes significant improvements in latency, throughput, and memory efficiency under realistic domain workloads, without modifying model computations. The methodology heralds a new direction for scalable, specialization-preserving MoE infrastructure, informing future research in distributed inference and training of large sparse networks.

Reference: "Least-Loaded Expert Parallelism: Load Balancing An Imbalanced Mixture-of-Experts" (2601.17111)