最小負荷エキスパート並列処理：不均衡な混合エキスパートの負荷分散

要旨

Mixture-of-Experts（MoE）モデルは、統計的に均衡した専門家ルーティングを保証するため、通常は明示的な負荷分散制約を課して事前学習されます。しかし、十分に学習されたMoEモデルであっても、著しく不均衡なルーティングが生じることが観察されています。この振る舞いは、ある意味で自然であり、むしろ望ましいものと言えます。なぜなら、不均衡なルーティングにより、モデルはドメイン固有の知識を専門家のサブセット内に集中させることができるからです。専門家並列処理（EP）は、専門家を複数のデバイスに分散させることでMoEモデルのスケーリングを図るように設計されていますが、その前提としてあまり議論されていないのが、均衡したルーティングです。極端な不均衡下では、EPは不均衡に多数のトークンを少数の専門家に集中させてしまい、事後学習や推論時に過負荷デバイス上で計算資源やメモリ資源の限界による障害を引き起こす可能性があります。この段階では明示的な負荷分散は往往にして適用不可能です。本研究では、Least-Loaded Expert Parallelism（LLEP）という新たなEPアルゴリズムを提案します。LLEPは、過負荷デバイスから余剰トークンおよび関連する専門家パラメータを未使用のデバイスに動的に再ルーティングします。これにより、メモリ制約を遵守しつつ、全てのデバイスが最小の全体レイテンシでワークロードを完了することが保証されます。様々なモデル規模において、LLEPは標準的なEPと比較して最大5倍の高速化と、ピークメモリ使用量の4倍の削減を実現しました。これにより、より高速かつ高スループットな事後学習と推論が可能となり、gpt-oss-120bでは約1.9倍の高速化を達成しました。本手法は、詳細な理論分析と、 ablation study を含む包括的な実証評価によって裏付けられています。これらの結果は重要なトレードオフを明らかにし、特定のハードウェアにおいて最適な性能を達成するためのハイパーパラメータチューニングの原理的な枠組みを可能にします。

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.

最小負荷エキスパート並列処理：不均衡な混合エキスパートの負荷分散

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

要旨

Support