AI・機械学習
中級

MoE Load Balancing(エムオーイー ロードバランシング)

Sparse MoE モデルの学習・推論時にエキスパート間のトークン割り当てを均等化する手法群の総称。負荷の偏りはエキスパート崩壊や計算資源の無駄遣いを引き起こすため、補助損失やキャパシティ制約で制御する。

0 回閲覧
0 いいね

MoE Load Balancing が必要な理由

Sparse MoE モデルでは、ルーターが各トークンを少数のエキスパートに割り当てます。しかし自然なデータ分布のもとでは、特定のエキスパートに人気が偏りやすく、放置すると深刻な問題が発生します。

学習面の問題: 過剰に選択されるエキスパートは過学習しやすく、逆にほとんど選択されないエキスパートは学習が進まず死重化します。結果としてモデルの実効容量が縮小し、Dense モデルに対する優位性を失います。

推論面の問題: 分散推論環境では各エキスパートが異なる GPU に配置されます。負荷が偏ると特定の GPU がボトルネックとなり、他の GPU がアイドル状態で待機する非効率が生じます。

負荷分散手法の体系

主要な負荷分散手法を体系的に整理すると、大きく 3 つのアプローチに分類できます。

アプローチ手法例特徴
補助損失(Auxiliary Loss)Switch Transformer の負荷分散損失訓練目標に均一化ペナルティを追加
キャパシティ制約Expert Buffer / Token Droppingエキスパートの受入上限を設定
逆方向ルーティングExpert Choiceエキスパートがトークンを選ぶ逆転発想

Auxiliary Load Balancing Loss

最も広く採用される手法で、Switch Transformer で提案されました。損失関数は以下の形式です。

L_aux = α × N × Σ_i(f_i × P_i)

ここで f_i はエキスパート i に実際にルーティングされたトークンの割合、P_i はエキスパート i のルーティング確率の平均、N はエキスパート数、α は損失係数です。

f_i と P_i の両方が均一なとき L_aux が最小化されるため、学習を通じて均等な割り当てが促進されます。α の調整が重要で、大きすぎると専門化を阻害し、小さすぎると負荷偏りを許容します。一般的に α = 0.01〜0.1 の範囲で設定されます。

Token Dropping と Expert Buffer

各エキスパートにキャパシティファクター C を設定し、1 バッチ内で受け入れるトークン数の上限を C × (T/N) とします(T は総トークン数、N はエキスパート数)。上限を超えたトークンは Dropout されるか、残差接続経由で直接出力されます。

GShard では C = 2.0 が標準的な値とされ、平均の 2 倍までの偏りを許容します。C が小さいほど均一性は高まりますが、重要なトークンが Dropping される可能性が増えます。

Expert Choice Routing

従来の Top-K Routing がトークン側からエキスパートを選ぶのに対し、Expert Choice では各エキスパートが全トークンから上位 K 個を選択します。これにより各エキスパートの処理トークン数が完全に均一化され、補助損失なしで完璧な負荷分散が実現します。

ただし、同一トークンが複数エキスパートに選択されたり、どのエキスパートにも選択されないトークンが発生する可能性があるため、自己回帰的な生成タスクでは使いにくいという制約があります。

DeepSeek-V3 のロスフリー負荷分散

DeepSeek-V3 は補助損失を使わずにバイアス項の調整だけで負荷分散を実現する手法を提案しました。各エキスパートにバイアス b_i を持たせ、負荷が多いエキスパートのバイアスを下げ、少ないエキスパートのバイアスを上げることで動的に均一化します。補助損失によるメイン損失への干渉を排除できるため、モデル品質の劣化なく負荷分散が可能です。

分散推論での負荷分散

学習時だけでなく推論時の負荷分散も重要です。Expert Parallelism(EP)ではエキスパートを異なる GPU に配置し、All-to-All 通信でトークンを適切な GPU に転送します。負荷が偏ると特定 GPU の処理待ちが発生するため、推論フレームワーク(vLLM, TensorRT-LLM 等)では動的なトークンバッチングやエキスパート複製で対応します。

FAQ

Q1: 負荷分散は推論品質に影響しますか?

A1: 強すぎる負荷分散(α が大きい補助損失やキャパシティファクターが小さい設定)はエキスパートの専門化を阻害し、品質低下につながります。DeepSeek-V3 のロスフリー手法はこのトレードオフを解消する有力なアプローチです。

Q2: 負荷の偏りをリアルタイムでモニタリングできますか?

A2: はい。各エキスパートのトークン処理数、ルーティング確率の分散、最大/最小負荷比などをバッチごとに記録するのが一般的です。WandB や TensorBoard でダッシュボード化すると異常検知が容易になります。

Q3: エキスパート数が増えると負荷分散は難しくなりますか?

A3: エキスパート数が増えると各エキスパートあたりのトークン数が減少し、統計的な揺らぎが大きくなるため、負荷分散はより困難になります。ファイングレインド分割(DeepSeek-V2/V3)は多数エキスパートでも安定した負荷分散を実現する設計です。