詳細に潜む悪魔：トレーニング用の負荷分散損失の実装について 専門家モデルの混合に焦点を当てて

Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models

January 21, 2025

著者: Zihan Qiu, Zeyu Huang, Bo Zheng, Kaiyue Wen, Zekun Wang, Rui Men, Ivan Titov, Dayiheng Liu, Jingren Zhou, Junyang Lin

cs.AI

要旨

この論文では、Mixture-of-Experts（MoEs）モデルのトレーニング時にLoad-balancing Loss（LBL）の実装を再検討します。具体的には、MoEsのためのLBLは、N_Eがエキスパートの総数であり、f_iがエキスパートiが選択される頻度を表し、p_iがエキスパートiの平均ゲーティングスコアを示すと定義されます。既存のMoEトレーニングフレームワークは通常、並列トレーニング戦略を採用しており、f_iとLBLがマイクロバッチ内で計算され、それから並列グループ全体で平均化されます。本質的に、10億スケールのLLMsのトレーニング用のマイクロバッチには通常、非常に少数のシーケンスが含まれています。そのため、マイクロバッチのLBLはほぼシーケンスレベルであり、ルーターは各シーケンス内でトークンを均等に分配するようにプッシュされます。この厳格な制約の下では、ドメイン固有のシーケンス（例：コード）からのトークンさえもすべてのエキスパートに均等にルーティングされ、エキスパートの特殊化が抑制されます。本研究では、この制約を緩和するためにグローバルバッチを使用してLBLを計算することを提案しています。グローバルバッチにはマイクロバッチよりも多様なシーケンスが含まれるため、コーパスレベルでの負荷分散が促進されます。具体的には、エキスパート間のf_iをマイクロバッチ間で同期させるための追加の通信ステップを導入し、それを使用してLBLを計算します。MoEsベースのLLMs（合計42.8Bのパラメータと400Bのトークンまで）のトレーニング実験を通じて、グローバルバッチLBL戦略が事前トレーニングのパープレキシティとダウンストリームタスクの両方で優れたパフォーマンス向上をもたらすことが驚くほど明らかになりました。当社の分析によれば、グローバルバッチLBLはMoEエキスパートのドメイン特殊化も大幅に改善します。

English

This paper revisits the implementation of Load-balancing Loss (LBL) when training Mixture-of-Experts (MoEs) models. Specifically, LBL for MoEs is defined as N_E sum_{i=1}^{N_E} f_i p_i, where N_E is the total number of experts, f_i represents the frequency of expert i being selected, and p_i denotes the average gating score of the expert i. Existing MoE training frameworks usually employ the parallel training strategy so that f_i and the LBL are calculated within a micro-batch and then averaged across parallel groups. In essence, a micro-batch for training billion-scale LLMs normally contains very few sequences. So, the micro-batch LBL is almost at the sequence level, and the router is pushed to distribute the token evenly within each sequence. Under this strict constraint, even tokens from a domain-specific sequence (e.g., code) are uniformly routed to all experts, thereby inhibiting expert specialization. In this work, we propose calculating LBL using a global-batch to loose this constraint. Because a global-batch contains much more diverse sequences than a micro-batch, which will encourage load balance at the corpus level. Specifically, we introduce an extra communication step to synchronize f_i across micro-batches and then use it to calculate the LBL. Through experiments on training MoEs-based LLMs (up to 42.8B total parameters and 400B tokens), we surprisingly find that the global-batch LBL strategy yields excellent performance gains in both pre-training perplexity and downstream tasks. Our analysis reveals that the global-batch LBL also greatly improves the domain specialization of MoE experts.