Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts

要約

Mixture-of-Experts (MoE) モデルの場合、不均衡なエキスパートの負荷は、ルーティングの崩壊や計算オーバーヘッドの増加につながります。
既存の方法では負荷バランスを促進するために補助損失が一般的に採用されていますが、補助損失が大きいとトレーニングに無視できない干渉勾配が導入され、モデルのパフォーマンスが低下します。
トレーニング中に望ましくない勾配を生成しないように負荷バランスを制御するために、補助損失のない負荷分散戦略を特徴とする損失のない分散を提案します。
具体的には、トップ K のルーティングを決定する前に、ロスフリー バランシングはまず各エキスパートのルーティング スコアにエキスパートごとのバイアスを適用します。
最近の負荷に応じて各エキスパートのバイアスを動的に更新することにより、ロスフリー バランシングはエキスパート負荷のバランスのとれた分散を一貫して維持できます。
さらに、ロスフリー バランシングは干渉勾配を生成しないため、MoE トレーニングから得られるモデル パフォーマンスの上限も引き上げられます。
最大 200B のトークンでトレーニングされた最大 3B のパラメーターを使用して、MoE モデルでのロスフリー バランシングのパフォーマンスを検証します。
実験結果によると、ロスフリー バランシングは、従来の補助損失制御のロード バランシング戦略と比較して、優れたパフォーマンスと優れたロード バランシングの両方を実現します。

要約(オリジナル)

For Mixture-of-Experts (MoE) models, an unbalanced expert load will lead to routing collapse or increased computational overhead. Existing methods commonly employ an auxiliary loss to encourage load balance, but a large auxiliary loss will introduce non-negligible interference gradients into training and thus impair the model performance. In order to control load balance while not producing undesired gradients during training, we propose Loss-Free Balancing, featured by an auxiliary-loss-free load balancing strategy. To be specific, before the top-K routing decision, Loss-Free Balancing will first apply an expert-wise bias to the routing scores of each expert. By dynamically updating the bias of each expert according to its recent load, Loss-Free Balancing can consistently maintain a balanced distribution of expert load. In addition, since Loss-Free Balancing does not produce any interference gradients, it also elevates the upper bound of model performance gained from MoE training. We validate the performance of Loss-Free Balancing on MoE models with up to 3B parameters trained on up to 200B tokens. Experimental results show that Loss-Free Balancing achieves both better performance and better load balance compared with traditional auxiliary-loss-controlled load balancing strategies.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google