Expert Parallelism

エキスパート並列化（EP）は、分散学習と推論に用いる、Mixture of Experts（MoE）モデル向けの戦略です。異なるエキスパートを別々のGPUに配置し、各tokenを、そのトークンに選ばれたエキスパートを保持するデバイスへ送ります。MoEレイヤーでは、Transformerの大きなフィードフォワードブロックを、多数の小さなエキスパートネットワークとルーターに置き換えます。ルーターは各トークンに対して一部のエキスパートだけを有効にするため、トークンごとの計算量よりも総パラメーター数をはるかに速く増やせます。エキスパート並列化を使うと、パラメーターをメモリに収めやすくなります。すべてのGPUがすべてのエキスパートを保持するのではなく、それぞれが一部を保持するためです。GoogleのNoam Shazeer氏らは、2017年のスパースゲートMoEレイヤーで、共有コピーを1つずつ各デバイスに分散して配置していました。GShardは2020年6月にこの手法をスケールさせ、6,000億を超えるパラメーターを持つ翻訳モデルのエキスパートを、2,048台のTPU v3アクセラレーターに分散しました。現在、エキスパート並列化は大規模なMoE 言語モデルや視覚言語モデルの学習とサービングに広く使われています。

トークンがGPU間を移動する仕組み#

各MoE層は、次の4つのステップで動作します。

1. ルーティング: ルーターが各トークンをすべてのエキスパートに対してスコアリングし、一般的には1つ、2つ、または8つの上位k個のエキスパートを選びます。
2. ディスパッチ： All-to-all型の集合通信を使って、各トークンの隠れ状態を、選択されたエキスパートを保持するGPUへ送ります。
3. 計算: 各GPUは、受け取ったトークンに対してローカルのエキスパートを実行します。
4. 集約: 2回目のall-to-all通信で出力を各トークンの元のGPUに戻し、そこでルーターのスコアに基づいて重み付けし、合計します。

以下のコードスニペットでは、8個のエキスパートを4台のGPUに分散し、GPUごとに2個ずつ配置した1つの層について、ルーティングとディスパッチのステップをシミュレートします。

import torch

torch.manual_seed(0)
num_experts, num_gpus, top_k = 8, 4, 2
tokens = torch.randn(4096, 512)  # 4,096個のトークンの隠れ状態
router = torch.nn.Linear(512, num_experts)

# ルーティング: 各トークンはルーターのスコアに基づいて上位2つのエキスパートを選択します
with torch.no_grad():
    _, experts = router(tokens).topk(top_k, dim=-1)

# エキスパート並列化: エキスパート0～1はGPU 0、エキスパート2～3はGPU 1に配置され、以降も同様です
gpu = experts // (num_experts // num_gpus)

# ディスパッチ: 各エキスパートとGPUが受け取るトークンとエキスパートの割り当て数を数えます
print(torch.bincount(experts.flatten(), minlength=num_experts).tolist())
print(torch.bincount(gpu.flatten(), minlength=num_gpus).tolist())

出力の1行目は各エキスパートの負荷を示します。割り当て数は948～1,087で、均等に分配した場合の1,024に対する値です。2行目は、8,192件の割り当てのうち、4台のGPUがそれぞれ1,968～2,139件を受け取ったことを示します。ランダムに初期化されたルーターでさえ処理を不均等に分散します。また、学習済みのルーターでは人気の高いエキスパートが生じ、負荷の偏りがさらに大きくなります。

負荷分散と容量#

層の処理は、最も負荷の高いGPUが完了して初めて終わるため、ルーティングの偏りはハードウェアのアイドル状態に直結します。MoEの学習では、次の3つの方法でこれを抑えます。

• 補助負荷分散損失： 損失関数に追加する項で、トークンがエキスパート間に均等に分散されるようルーターに促します。
• 容量係数： 各エキスパートが1バッチで受け付けるトークン数の上限を固定します。たとえば均等に分配した数の1.25倍を上限とし、超過したトークンはそのエキスパートをスキップします。
• バイアスベースの負荷分散： DeepSeek-V3は、ルータースコアにエキスパートごとのバイアスを加えます。このバイアスは選択するエキスパートにのみ影響します。各ステップの後、過負荷のエキスパートではバイアスを下げ、稼働の少ないエキスパートでは上げることで、主目的関数と競合する補助損失なしに負荷を分散します。

容量の上限がある場合、各トークンの経路はエキスパートの混雑状況によって変わります。

スキップされたトークンは、残差接続を通って変更されずにレイヤーを進むため、多くのトークンを破棄すると品質が低下します。DeepSeek-V3はバイアスベースの負荷分散に頼り、学習中にトークンを一切破棄しません。

もう一つのコストは通信です。MoEレイヤーごとに2回のall-to-all通信が発生します。そのため、エキスパート並列化は、GPUが高速なNVLink接続を共有するノード内で最も効率よく動作しますが、ノード間の低速な接続ではコストが高くなります。DeepSeek-V3は、8ノードにわたる64-wayエキスパート並列化を使って、各レイヤーのルーティング対象エキスパート256個を分散します。各トークンの送信先を最大4ノードに制限し、通信と計算を重ねて実行します。

並列化戦略の組み合わせ#

エキスパート並列化は、デバイス間で学習を分割する複数の方法の一つです。

• データ並列化： 各GPUにモデル全体を複製し、バッチを分割して、各ステップの後に勾配を同期します。
• テンソル並列化： 個々の重み行列をGPU間で分割するため、各GPUがすべてのレイヤーに参加します。
• パイプライン並列化： 連続するレイヤーを別々のGPUに割り当て、マイクロバッチを順に流します。
• コンテキスト並列化： 1つの長い入力シーケンスをGPU間で分割します。

1つのGPUに収まるDenseモデルはデータ並列化だけで十分です。これは、Ultralytics YOLO26が、device=[0, 1]のようなデバイスのリストを渡してトレーニングモードで複数のGPUを使って学習する方法です。大規模なMoEモデルでは複数の方式を組み合わせます。DeepSeek-V3は、6710億個のパラメーター（トークンごとのアクティブ数は370億個）を、16分割のパイプライン並列化、64分割のエキスパート並列化、シャーディングされたデータ並列化を用いて学習し、テンソル並列化は使用しませんでした。推論エンジンのvLLMなども、モデルをサービングする際に同じ分割を適用します。そのため、エキスパートが1つのGPUのメモリに収まらないMoEモデルも複数のGPUで実行でき、各GPUはエキスパートの一部だけを保持します。