[论文解读] Least-Loaded Expert Parallelism: Load Balancing An Imbalanced Mixture-of-Experts

Xuan-Phi Nguyen, Shrey Pandit+3 more|arXiv (Cornell University)|Jan 23, 2026

Mobile Crowdsensing and Crowdsourcing被引用 0

一句话总结

论文提出 Least-Loaded Expert Parallelism (LLEP)，一种动态路由方案，将超载 GPU 上的多余 token 与专家参数重新分配给资源利用不足的 GPU，以在失衡的 MoE 模型中实现负载平衡，相较于标准 Expert Parallelism (EP)，达到显著的加速和内存节省。

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.

研究动机与目标

• 动机：在训练好的 MoE 模型中，专家路由天然不平衡的问题及其对 EP 效率的影响。
• 提出 LLEP，在遵守内存约束的同时，动态将超出负载的 token 和专家权重重新路由到负载较低的设备。
• 对不平衡下的延迟和内存进行理论与经验分析，并在实际模型上展示实际收益。
• 提供面向硬件的调优指南，在后训练和推理场景中最大化吞吐量。

提出的方法

• 定义问题：在后训练或推理期间，EP 下 MoE 层的 token 路由不平衡。
• 提出 LLEP，使用 least-loaded assignment (LLA) 算法触发超出负载的 token 从过载 GPU 溢出到负载较低的 GPU。
• 开发一个溢出例程 (LLAS)，在 GPU 之间传输剩余工作负载及相应权重。
• 展示完整的 LLEP dispatch-combine 工作流程，包括向后传播支持和精确 MoE 计算。
• 提供延迟和峰值内存分析，以证明何时以及如何发生溢出，并引入可调的硬件感知参数 (α, m, λ)。
• 展示跨多种 MoE 架构的端到端和受控实验，显示加速和内存降低。

实验结果

研究问题

• RQ1 在预训练、微调或推理过程中，最前沿的 MoE 模型中不平衡路由如何表现？
• RQ2 负载感知的分布式路由策略是否能在不改变 MoE 行为的前提下减少每个 GPU 的延迟和峰值内存？
• RQ3 在不平衡条件下，使用和不使用 least-loaded 路由的 MoE 的理论与经验成本动态？
• RQ4 超参数 α、m、λ 如何影响 LLEP 在不同模型规模和硬件配置下的性能？
• RQ5 端到端部署（例如 gpt-oss-20b/120b）在吞吐量和内存稳定性方面是否比标准 EP 从 LLEP 中受益？

主要发现

• LLEP 在极端不平衡下相对于标准 EP 可实现高达 5× 的加速，同时保持内存使用稳定。
• LLEP 的每 GPU 峰值内存在不同不平衡场景下几乎保持不变，而标准 EP 可增长多达 4×。
• 在真实模型中，端到端吞吐量提升在 gpt-oss-20b 最高可达 2.2×，在 gpt-oss-120b 为 1.9×。
• 在实际开销下，使用 LLEP 训练比 EP 收敛速度快约 1.25×。
• 消融研究显示更大的批量带来更大加速；更高的 α 会降低加速，表明在大规模下倾向于平衡工作负载。