# エキスパート並列

[Expert parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=expert_parallelism) is a parallelism strategy for [mixture-of-experts (MoE) models](https://huggingface.co/blog/moe). 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](/docs/transformers/main/ja/expert_parallelism#transformers.DistributedConfig) クラスと `ep_size` 引数を使用して、エキスパート並列化を有効にします。ほとんどの MoE モデルはデフォルトで [トークンディスパッチ](#token-dispatch) を使用するため、`ep_size` は `tp_size` に依存せずに設定できます。Llama 4 や Gemma 4 のように、デフォルトでマスキングとオールリデュース（`model.ep_plan` における `"ep_router"` および `"moe_tp_experts"`）を使用するモデルもいくつかあります。マスキングは、`ep_plan` のオーバーライドを持つ任意のモデルでも利用可能ですが、エキスパートグループ内のすべてのランクが同じトークンを受け取るようにするため、`ep_size=tp_size` が必要です。

```py
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` の場合、両方が適用されます: [テンソル並列](./perf_infer_gpu_multi) プランはアテンションと稠密な 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` 内のエキスパートフォワードルールは、トークンがエキスパートにどのように到達するかを選択します。

| ルール | メカニズム | レイアウト |
| :--- | :--- | :--- |
| ルーター上で`"moe_tp_experts"`を`"ep_router"`と併用する | masking と all-reduce: 各 rank がバッチ全体に対してローカルのエキスパートを実行し、ルーターがそれらをマスクし、all-reduce によって出力を結合します | `ep_size=tp_size` |
| `"ep_dispatch_experts"` | [トークンディスパッチ](#token-dispatch): 各ランクは自身のトークンを保持し、ルーティングされた (token, expert) ペアを2回の all-to-all コレクティブ通信で交換するだけです | `ep_size` は `tp_size` の倍数であり、`fsdp_size * tp_size` を割り切れる値 |

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

[torchrun](https://pytorch.org/docs/stable/elastic/run.html)で推論スクリプトを起動します。プロセス数は`tp_size * fsdp_size * pp_size`と等しくする必要があり、`ep_size`はエキスパートの数を割り切れる必要があります。

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

### 計画の上書き

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

```py
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` などです）：

```py
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.

```py
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](./fsdp)によってシャード化され、それらの勾配が削減されます。
- エキスパートは`ep`にシャード化され、`efsdp_size = fsdp_size * tp_size // ep_size`が1より大きい場合は、さらに`efsdp`にFSDPシャード化されます。常にFSDPでラップされるため、`fsdp_mixed_precision`および`fsdp_cpu_offload`も適用され、[save_pretrained()](/docs/transformers/main/ja/main_classes/model#transformers.PreTrainedModel.save_pretrained)は他のパラメータと同様にこれらを収集します。
- Expert parallel group は `ep_size / tp_size` 個のバッチを保持するため、ある expert の勾配は、その数のバッチに対する和になります。`efsdp` の reduction は、その group のサイズではなく `fsdp_size` で割るため、dense モジュールに対して FSDP2 が採用するバッチごとの平均値と同じになります。
- すべてのローカルエキスパートは、レイヤーごとにゼロパッド行を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つのエキスパート:

```text
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 プロセスの場合:

```py
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 を使用します。

| 設定 | 結果 |
| :--- | :--- |
| `DistributedConfig(tp_size=4, fsdp_size=2, ep_size=4)` | 4つのTPグループでディスパッチし、各グループがバッチを4分割します(デフォルトのプラン)。レガシーのマスク付きEPを使用するには、ep_planを指定してください。 |
| `DistributedConfig(tp_size=1, fsdp_size=8, ep_size=4)` | 各ランクで独立したバッチをディスパッチし、スライシングを行わず、エキスパートをランクのペア間でFSDPシャードします。 |
| `DistributedConfig(tp_size=2, fsdp_size=4, ep_size=4)` | Dispatch はバッチごとに TP ペアを割り当て、各ペアが自分のバッチを2分割します。エキスパートグループごとに2つのバッチです。 |
| `DistributedConfig(tp_size=8, ep_size=8)` | 各ランクが1つのバッチを共有し、それを8分割してディスパッチするか、マスク付きプランのためにマスクとオールリデュースを行う。 |

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

## FSDP2との組み合わせ

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

```py
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](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer)で学習します。両方のメッシュにわたる勾配ノルムを取得し、各メッシュに独自のオプティマイザーのパラメータグループを割り当てます。[save_model()](/docs/transformers/main/ja/main_classes/trainer#transformers.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 はやや低下します。

| 設定 | tokens/s/GPU | ピークメモリ/GPU |
|---|---|---|
| `tp_size=8` | 3485 | 38.6 GB |
| `tp_size=4, fsdp_size=2` | 2900 | 34.2 GB |
| `tp_size=2, fsdp_size=4` | 2830 | 32.3 GB |

> [!WARNING]
> ロード時にシャード化されたモデルでは、チェックポイントからの再開はまだサポートされていないため、[Trainer](/docs/transformers/main/ja/main_classes/trainer#transformers.Trainer)はそれらに対して`save_only_model=True`または`save_strategy="no"`のみを受け入れます。

## API リファレンス[[transformers.DistributedConfig]]

#### transformers.DistributedConfig[[transformers.DistributedConfig]]

```python
transformers.DistributedConfig(tp_size: int | None = None, tp_plan: typing.Union[dict[str, str], typing.Literal['auto'], NoneType] = None, enable_sequence_parallel: bool = False, enable_expert_parallel: bool = False, fsdp_size: int | None = None, fsdp_plan: dict[str, str] | None = None, fsdp_cpu_offload: bool = False, fsdp_mixed_precision: bool = False, pp_size: int | None = None, ep_size: int | None = None, ep_plan: dict[str, str] | None = None)
```

[ソース](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/configuration_utils.py#L26)

**パラメータ:**

tp_size (`int`, *任意*) : テンソル並列化に使用するデバイス数。`None` かつ `tp_plan` が設定されている場合、デフォルト値は `WORLD_SIZE // (other_parallel_size)` になります。`None` かつ `tp_plan` が設定されていない場合、デフォルト値は 1 になります。

tp_plan (`dict[str, str]` または `"auto"`、*省略可*) : テンソル並列シャーディング計画。`tp_size` が設定されている場合、`"auto"` を渡すか `None` のままにすると、モデルで事前定義された `base_model_tp_plan` が使用されます。辞書を渡すことで、その計画の個別ルールを上書きできます。指定されていないルールは保持されます。空の辞書を渡すと、テンソル並列ルールは適用されません。

enable_sequence_parallel (`bool`、*任意*、デフォルトは`False`) : シーケンス並列処理用に予約されています。まだ実装されていません。

enable_expert_parallel (`bool`, *任意*、デフォルトは `False`) : 非推奨です。v5.21 で削除されました。`ep_size` を指定しない場合、`ep_size=tp_size, tp_plan={}` と同じです。`tp_size` 個のランク全体で、エキスパート並列プランのみを適用します。明示的に `ep_size` を指定した場合は、そちらが優先されます。このフラグは `tp_size` や `fsdp_size` を変更しません。

fsdp_size (`int`, *optional*) : FSDP（データ並列）に使用するデバイス数。`None` かつ `tp_size` が設定されている場合、デフォルト値は 1 です。

fsdp_plan (`dict[str, str]`、*省略可*) : FSDP2 のシャーディング計画。`None` の場合、モデルで定義済みの `base_model_fsdp_plan` とクラスレベルの `_fsdp_plan` が使用されます。辞書を渡すことで個別のルールを上書きできます。指定されていないルールは保持されます。`fsdp_size > 1` の場合のみ適用されます。

fsdp_cpu_offload (`bool`、*任意*、デフォルトは`False`) : FSDP2 の CPU オフロードを有効にするかどうか。

fsdp_mixed_precision (`bool`, *任意*, デフォルトは `False`) : FSDP2 で混合精度を有効にするかどうか。

pp_size (`int`, *optional*) : パイプライン並列化に使用するデバイスの数。`None`で他の並列モードが設定されている場合、デフォルトで1になります。

ep_size (`int`, *任意*) : 個別のエキスパートシャードを所有するデバイスの数です。デフォルト値は 1 です。EP を有効にするには、明示的に設定してください。`tp_size` の倍数であり、かつ `fsdp_size * tp_size` で割り切れる必要があります。エキスパートプランの All-reduce には `ep_size=tp_size` が必要ですが、トークンディスパッチ (`"ep_dispatch_experts"`) では `ep_size > tp_size` も許可されています。

ep_plan (`dict[str, str]`、*省略可*) : Expert parallel sharding plan。モデルの事前定義済みの`base_model_ep_plan`を使用するには`None`のままにしてください。辞書を渡すことで、そのプランの個別ルールを上書きできます。指定されていないルールは保持されます。`ep_size > 1`の場合にのみ適用され、同じモジュールに対しては`tp_plan`のルールよりも優先されます。`"ep_dispatch_experts"`ルールは、ルーターのマスクングとall-reduceの代わりに、all-to-allトークンディスパッチを選択します。

テンソル並列、パイプライン並列、またはFSDP2並列によるネイティブ分散推論および学習の設定。