解决专家模型负载不均导致的计算瓶颈问题,提升推理效率。
Least-Loaded Expert Parallelism: Load Balancing An Imbalanced Mixture-of-Experts
- 动态迁移过载设备上的专家参数与令牌,实现负载均衡。
- 相比传统方法,最高提速5倍,峰值内存减少4倍。
- 适合大规模MoE模型的高效训练与推理部署。
混合专家(MoE)模型通常通过显式负载均衡约束进行预训练,以确保专家路由的统计平衡。然而,我们发现即使训练良好的MoE模型仍存在显著的路由不平衡。这种现象在一定程度上是自然且有益的,因为它允许模型将特定领域知识集中于少数专家中。专家并行(EP)通过将专家分布在多个设备上来扩展MoE模型,但其隐含假设为路由平衡。在极端不平衡情况下,EP可能将大量令牌集中到少数专家,导致部分设备在后训练或推理阶段因计算和内存压力而失败,此时显式负载均衡难以应用。为此,我们提出最少负载专家并行(LLEP),一种新型的EP算法,可动态将过载设备上的多余令牌及对应专家参数重路由至空闲设备,确保所有设备在最小集体延迟内完成任务,同时满足内存限制。在不同模型规模下,LLEP相较标准EP最高实现5倍加速和4倍峰值内存降低,支持更快速、高吞吐的后训练与推理,例如gpt-oss-120b实现约1.9倍加速。我们通过详尽的理论分析与全面的实证评估(包括消融实验)验证了该方法,揭示关键权衡,并建立面向硬件特性的超参数调优框架,以实现最优性能。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。