Mixture-of-Experts向けの補助損失フリーな負荷分散戦略

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

August 28, 2024

著者: Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, Damai Dai

cs.AI

要旨

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

English

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.