エキスパート並列

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_dispatch_experts" でエキスパートをシャーディングします。その実装は、トークンがエキスパートにどのように到達するかを選択します。トークンディスパッチにより、ep_size と tp_size を異ならせることができます。一部の EP 実装では ep_size=tp_size が必要であり、例えば "all_reduce" や "deepgemm_megamoe" では、エキスパートグループのすべてのランクが同じトークンを保持します。

Llama 4にはEPプランがありません。そのエキスパートは、ルーターのスコアで既にスケーリング済みのすべてのトークンを処理するため、ディスパッチするものがありません。代わりに、テンソル並列性(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_dispatch_experts"ルールは、トークンがエキスパートにどのように到達するかを選択します。

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

"ep_dispatch_experts" は、ランクごとにエキスパート全体をシャード化します（ep_size は tp_size の倍数であり、fsdp_size * tp_size を割り切ります）。アクセスには ep_implementation で指定された実装を使用します。実装の選択を参照してください。

実装の選択

"ep_dispatch_experts" は、エキスパートを一度だけシャード化し、DistributedConfig の ep_implementation で選択される EPImplementationInterface に登録された実装から、トークンがエキスパートに到達する方法を決定します。

ep_implementationは、DistributedConfigで設定された並列サイズの中から選択するだけです。エキスパートグループのランクが異なるバッチを保持する場合、例えばtp_size=1かつfsdp_size=ep_sizeの場合、適用可能な実装は"all_to_all"のみです。

実装の追加

実装はEPImplementationをサブクラス化し、ep_implementationの名前で登録されます。各expertsモジュールごとに一度だけ構築されるため、通信バッファなどの状態を保持できます。このランクが保持するバッチと、グローバルなexpert IDを引数として呼び出すと、各トークンに対する重み付きexpert出力が返されます。このランクのexpertは、experts_forwardによって、行、ローカルexpert ID、およびそれらに付与された重みを用いて実行されます。

from transformers.distributed.expert_parallel import EPImplementation, EPImplementationInterface

class MyEP(EPImplementation):
    def __call__(self, hidden_states, top_k_index, top_k_weights):
        tokens, expert_ids, weights = ...  # bring this rank's experts the rows routed to them
        expert_output = self.experts_forward(tokens, expert_ids, weights)
        return ...  # the weighted expert output for each token of the batch this rank holds

EPImplementationInterface.register("my_ep", MyEP)
distributed_config = DistributedConfig(tp_size=1, fsdp_size=8, ep_size=8, ep_implementation="my_ep")

勾配は、トークンを保持するすべてのランク（ルーティング重みを含む）に到達する必要があります。そのため、autograd が微分しない集合通信を torch.autograd.Function でラップしてください。

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 を明示的に設定してください。

ep_size=tp_sizeの場合、デフォルトの"auto"実装は推論時に"all_reduce"を実行し、学習時にトークンをディスパッチします。学習時にもマスクとall-reduceを行うには、ep_implementation="all_reduce"を設定してください。

distributed_config = DistributedConfig(tp_size=4, ep_size=4, ep_implementation="all_reduce")

トークンディスパッチ

Token dispatch は "all_to_all" の実装です。各ランクはバッチの自分の担当部分で学習を行います。各 MoE レイヤーにおいて、ランクは自分のトークンをルーティングし、all-to-all によって各 (token, expert) ペアをその expert を所有するランクに送信します。次に、受信したデータに対してローカルの expert を実行し、2 回目の all-to-all で結果を受け取り、ルーティング重みと結合します。移動するのはルーティングされた活性化値と expert の出力のみであり、どのランクも自分が所有していないトークンの expert を計算することはありません。

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 と ep にわたって重みをシャードしますが、オプティマイザの状態とルールのないモジュールは、グループのすべてのランクにレプリケートされ続けます。これにより、学習できるモデルのサイズが制限されます。別のメッシュ次元で FSDP2 を fsdp_size を用いて追加してください。マスキングと all-reduce を使用する場合、ep_size=tp_size を維持し、ep_implementation="all_reduce" を渡してください。

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_implementation="all_reduce",  # "auto" would dispatch tokens in training, which the Trainer does not support yet
)
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 リファレンス