エキスパート並列

Expert parallelism is a parallelism strategy for mixture-of-experts (MoE) models. Each expert’s feedforward layer lives on a different hardware accelerator. A router dispatches tokens to the appropriate experts and gathers the results. This approach scales models to far larger parameter counts without increasing computation cost because each token activates only a few experts.

DistributedConfig

DistributedConfig クラスと ep_size 引数を使用して、エキスパート並列化を有効にします。ほとんどの MoE モデルはデフォルトで トークンディスパッチ を使用するため、ep_size は tp_size に依存せずに設定できます。Llama 4 や Gemma 4 のように、デフォルトでマスキングとオールリデュース（model.ep_plan における "ep_router" および "moe_tp_experts"）を使用するモデルもいくつかあります。マスキングは、ep_plan のオーバーライドを持つ任意のモデルでも利用可能ですが、エキスパートグループ内のすべてのランクが同じトークンを受け取るようにするため、ep_size=tp_size が必要です。

import os

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.distributed.configuration_utils import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=int(os.environ["WORLD_SIZE"]),
    ep_size=int(os.environ["WORLD_SIZE"]),
)

model = AutoModelForCausalLM.from_pretrained(
    "openai/gpt-oss-120b",
    distributed_config=distributed_config,
)

各 MoE モデルは、その設定内で 2 つのプランを定義します: 稠密モジュール用の base_model_tp_plan と、エキスパート用の base_model_ep_plan です。これらは読み込まれたモデル上で model.tp_plan および model.ep_plan として公開されます。tp_size > 1 かつ ep_size > 1 の場合、両方が適用されます: テンソル並列 プランはアテンションと稠密な MLP をシャーディングし、エキスパート並列プランはエキスパートをシャーディングします。同じモジュールに対しては、EP ルールが TP ルールよりも優先されるため、エキスパートの重みは EP プランによって一度だけシャーディングされます。EP プランでは、GroupedGemmParallel スタイルはエキスパート次元に沿ってエキスパートの重みを分割し、各ランクがローカルなエキスパートのみを読み込むようにします。また、ep_router は、エキスパートの出力を all-reduce で結合する前に、他のランク上に存在するエキスパートをマスクします。

tp_plan は tp_size > 1 の場合のみ適用され、ep_plan は ep_size > 1 の場合のみ適用されます。TP が有効で EP が無効の場合、エキスパートのルールを含む完全な TP プランが適用されます。

ep_plan 内のエキスパートフォワードルールは、トークンがエキスパートにどのように到達するかを選択します。

enable_expert_parallel=True を ep_size なしで指定することは非推奨となり、FutureWarning が出力されます。これは ep_size=tp_size, tp_plan={} と同様に、エキスパートのみをシャーディングします。

torchrunで推論スクリプトを起動します。プロセス数はtp_size * fsdp_size * pp_sizeと等しくする必要があり、ep_sizeはエキスパートの数を割り切れる必要があります。

torchrun --nproc-per-node 8 your_script.py

計画の上書き

tp_plan={...} または ep_plan={...} を DistributedConfig に渡すことで、事前定義されたプランの個別ルールを上書きできます。指定されていないルールは保持され、マージされたプランはモデルに保存されます。各キーは、モジュール、パラメータ、または既存のプランエントリと一致する必要があります。一致しない場合、シャーディングが実行される前に ValueError が発生します。パスは、読み込まれたモデルから見た完全なパスを使用してください。例えば、causal LM には model.layers.*、そのベースモデルには layers.* を使用します。

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,
    tp_plan={"model.layers.*.self_attn.q_proj": "colwise_rep"},
    ep_plan={"model.layers.*.mlp.experts.down_proj": "grouped_gemm"},
)

プランを指定しても並列サイズは推定されません。tp_size と ep_size を明示的に設定してください。

ほとんどの MoE モデルでは、デフォルトで "ep_dispatch_experts" が使用されます。代わりにマスキングと all-reduce を使用するには、ep_size=tp_size を設定し、ルーターとエキスパートの forward ルールを両方オーバーライドしてください（ルーターのモジュール名はモデルによって異なり、gpt-oss の場合は mlp.router などです）：

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,
    ep_plan={
        "model.layers.*.mlp.gate": "ep_router",
        "model.layers.*.mlp.experts": "moe_tp_experts",
    },
)

逆に、プランがマスキングを使用しているモデルのエキスパートのフォワードルールを"ep_dispatch_experts"でオーバーライドし、トークンディスパッチを使用します。この場合、ディスパッチでは各エキスパートの所有者を特定するためにグローバルなエキスパートIDが必要になるため、ルータールールは無視されます。

トークンディスパッチ

With token dispatch, each rank trains on its own part of the batch. At every MoE layer, a rank routes its tokens, sends each (token, expert) pair to the rank that owns the expert with an all-to-all, runs its local experts on what it receives, gets the results back with a second all-to-all and combines them with the routing weights. Only the routed activations and expert outputs travel, and no rank computes experts for tokens it does not own.

from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=1,
    fsdp_size=8,
    ep_size=4,
)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B", distributed_config=distributed_config)

tp_size=1の場合、ep_sizeはfsdp_sizeとエキスパートの数を割り切れる必要があります。また、アテンションはnum_key_value_headsによって制約されません。モデルの残りの部分については:

• エキスパート以外のパラメータは、fsdp上でFSDP2によってシャード化され、それらの勾配が削減されます。
• エキスパートはepにシャード化され、efsdp_size = fsdp_size * tp_size // ep_sizeが1より大きい場合は、さらにefsdpにFSDPシャード化されます。常にFSDPでラップされるため、fsdp_mixed_precisionおよびfsdp_cpu_offloadも適用され、save_pretrained()は他のパラメータと同様にこれらを収集します。
• Expert parallel グループは ep_size / tp_size 個のバッチを保持するため、ある expert の勾配は、その数のバッチに対する和になります。efsdp の reduction は、そのグループサイズではなく fsdp_size で割るため、FSDP2 が dense モジュールに対して採用するバッチごとの平均値と同じになります。
• すべてのローカルエキスパートは、レイヤーごとにゼロパッド行を1行処理します。エキスパートがトークンを1つも受け取らなかったランクであっても、reverse all-to-allおよびエキスパートの勾配リダクションに参加します。

サイズの組み合わせ方

• tp_size * fsdp_size はプロセス数を表します。ep_size はプロセス数を増やしません。エキスパートの重みのために同じランクを再グループ化するだけです。
• ep_sizeは、エキスパートのリストをep_size個のブロックに分割します。各ランクはnum_experts / ep_size個のエキスパートを計算し、ep_size個の連続するランクが1つの完全なセットを保持します。このランクの集合が、all-to-allが実行されるグループです。
• efsdp_size = fsdp_size * tp_size / ep_size は、エキスパートセットの完全なコピーが何個存在するかを示します。異なるコピー内で同じ位置にあるランクは、メモリ削減のためにそれらのエキスパートをシャード化し、勾配を平均化します。これは FSDP がデンスモジュールに対して行う処理と同様です。
• あるランクが保持するバッチは、fsdpのみによって決まります。連続するランクがTPグループを形成し、同じバッチを取得します。fsdp_size個のグループはそれぞれ異なるバッチを取得します。

8つのプロセス、tp_size=2, fsdp_size=4、8つのエキスパート:

rank         0      1       2      3       4      5       6      7
batch       [====B0====]   [====B1====]   [====B2====]   [====B3====]     one batch per TP pair
tp           0      1       0      1       0      1       0      1

ep_size=2   E0-3   E4-7    E0-3   E4-7    E0-3   E4-7    E0-3   E4-7     group = a TP pair,   efsdp_size=4
ep_size=4   E0E1   E2E3    E4E5   E6E7    E0E1   E2E3    E4E5   E6E7     group = two pairs,   efsdp_size=2
ep_size=8   E0     E1      E2     E3      E4     E5      E6     E7       group = all ranks,   efsdp_size=1

この図から2つの数値が導き出せます：

• EPグループ内では、各トークンは、そのバッチを保持するペアのランクごとに1回ずつ、合計tp_size回存在します。これはep_sizeには依存しません。
• EPグループはep_size / tp_size個の異なるバッチを保持します。これは、あるエキスパートの勾配が合計される対象のバッチ数です。

テンソル並列性を使用する場合

tp_size > 1 を設定すると、TP プランに従って dense モジュールをシャーディングし、エキスパートは dispatch を使用します。8 プロセスの場合:

distributed_config = DistributedConfig(
    tp_size=2,
    fsdp_size=4,
    ep_size=4,
)

TP ランクの各ペアは同じバッチを受信します。これは、テンソル並列性がペア内のアクティベーションを複製するためです。両方のランクがすべてのトークンをディスパッチした場合、あるエキスパートを所有するランクはすべてのトークンを2回受信し、2回計算することになり、その重み勾配が2倍になってしまいます。エキスパートは1つのランクに完全に存在するため、重複分を重みで分割することはできません。また、所有権は通常別のランクにあるため、マスキングのように所有権で解決することもできません。したがって、ペアは行を分割します。各 TP ランクはトークンの互いに素な 1 / tp_size をディスパッチし、結果は送信元のランクに返され、ゼロパディングされた半分に対するペア間の all-reduce により、次のレイヤーが期待する複製された出力が復元されます。分割は tp_size によるものであり、ep_size ではありません。これは、バッチを保持しているランクのみが送信できるためです。エキスパートグループは4つのランクにまたがり、各エキスパートは efsdp_size = 2 のランク間で FSDP によりシャード化されます。一方、トランクの FSDP グループは4つのランクにまたがります。モデルの通常の TP 制約（アテンションヘッドの割り切れることなど）は、デンスモジュールにも引き続き適用されます。トークンスライスは不均一または空になる場合があり、単一トークンデコーディング中も同様です。

Token dispatch はまだパイプライン並列化と組み合わせられません。 pp_size=1 を使用してください（未テスト） これらの設定はそれぞれ 8 つの GPU を使用します：

Trainerは、トークンディスパッチをまだ考慮していません。バッチおよびトークンのカウントは、TPグループのランクがバッチを共有し、fsdp_size個のデータ並列シャードが存在する、オールリデュースレイアウトを前提としています。Trainerによるディスパッチへの対応は、後続のリリースで提供される予定です。

FSDP2との組み合わせ

Tensor parallelism と expert parallelism は tp にわたって重みをシャードしますが、オプティマイザの状態とルールのないモジュールは、グループのすべてのランクにレプリケートされたままです。そのため、学習できるモデルのサイズが制限されます。fsdp_size を用いて、第 2 のメッシュ次元に FSDP2 を追加します。マスキングと all-reduce を使用し、ep_size=tp_size を維持したまま、`ep_plan={“layers..mlp.gate”: “ep_router”, “layers..mlp.experts”: “moe_tp_experts”}“` を渡します。

from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,  # expert parallel size, must match tp_size with masking and all-reduce
    fsdp_size=2,  # data parallel shards
    ep_plan={
        "model.layers.*.mlp.gate": "ep_router",
        "model.layers.*.mlp.experts": "moe_tp_experts",
    },
)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B", distributed_config=distributed_config)

モデルを通常どおり読み込み、Trainerで学習します。両方のメッシュにわたる勾配ノルムを取得し、各メッシュに独自のオプティマイザーのパラメータグループを割り当てます。save_model()は、シャード化された重みを通常のチェックポイントに集約します。これにはaccelerate>=1.12が必要であり、Trainerがtp_sizeとfsdp_sizeを~Accelerate.ParallelismConfigに反映できるようにします。

下表は、8xH100 GPU 環境における EP のみの学習と 2D EP+FSDP2 の比較を示しています。ワークロードは、シーケンス長 2048 で bf16 による Qwen3-30B-A3B のフルファインチューニングです。FSDP シャード数を増やすとピークメモリは削減されますが、fsdp 間でエキスパートを all-gather および reduce-scatter するため、tokens/s はやや低下します。

ロード時にシャード化されたモデルでは、チェックポイントからの再開はまだサポートされていないため、Trainerはそれらに対してsave_only_model=Trueまたはsave_strategy="no"のみを受け入れます。

API リファレンス