Binary-Integer-Programming Based Algorithm for Expert Load Balancing in Mixture-of-Experts Models

要約

MOE（Expertsの混合）モデルの事前トレーニングの場合、主な問題の1つは不均衡な専門家負荷であり、ルーティングの崩壊または計算オーバーヘッドの増加を引き起こす可能性があります。
既存の方法には、損失制御された方法と損失のない方法が含まれています。この方法では、最初は不均衡な程度の両方のトレーニングステップが依然として高く、ゆっくりと減少します。
この作業では、バイナリ整数プログラミング（BIP）に基づいた専門家の負荷分散アルゴリズムであるBIPベースのバランスを提案します。
アルゴリズムは、各MOE層に追加のベクトルQを維持し、非常に小さな時間コストでバイナリ整数プログラミングを解くことにより、SのトップK順序を変更するのに役立ちます。
2つのMOE言語モデルにアルゴリズムを実装します：16-Expert（0.3b）および64-Expert（1.1b）。
実験結果は、両方のモデルが損失制御方法と損失のない方法と比較して、アルゴリズムが最も低い困惑でモデルを訓練し、損失制御方法と比較してトレーニング前の時間の少なくとも13％を節約することを示しています。
現在の知識の中で、これは、トレーニング前のプロセス全体で最初のステップから最後のステップまで、すべてのMOE層のすべての専門家の負荷バランスステータスを維持する最初のルーティングアルゴリズムであり、訓練されたMOEモデルもうまく機能します。
この作業のコード資料は、https：//github.com/sunyuanllm/bip_routing_algorithmで入手できます。

要約(オリジナル)

For pre-training of MoE (Mixture-of-Experts) models, one of the main issues is unbalanced expert loads, which may cause routing collapse or increased computational overhead. Existing methods contain the Loss-Controlled method and the Loss-Free method, where both the unbalanced degrees at first several training steps are still high and decrease slowly. In this work, we propose BIP-Based Balancing, an expert load balancing algorithm based on binary integer programming (BIP). The algorithm maintains an additional vector q on each MoE layer that can help change the top-K order of s by solving a binary integer programming with very small time costs. We implement the algorithm on two MoE language models: 16-expert (0.3B) and 64-expert (1.1B). The experimental results show that on both models comparing with the Loss-Controlled method and the Loss-Free method, our algorithm trains models with the lowest perplexities, while saves at least 13% of pre-training time compared with the Loss-Controlled method. Within our current knowledge, this is the first routing algorithm that achieves maintaining load balance status on every expert in every MoE layer from the first step to the last step during the whole pre-training process, while the trained MoE models also perform well. The code material of this work is available at https://github.com/sunyuanLLM/bip_routing_algorithm.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google