# 专家并行

[专家并行](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=expert_parallelism)是一种针对[混合专家（MoE）模型](https://huggingface.co/blog/moe)的并行策略。每个专家的 feedforward 层位于不同的硬件加速器上。路由器将 token 分派给相应的专家并收集结果。由于每个 token 仅激活少数几个专家，该方法可以在不增加计算成本的情况下，将模型扩展到远更大的参数量。

## DistributedConfig

使用 [DistributedConfig](/docs/transformers/main/zh/expert_parallelism#transformers.DistributedConfig) 类和 `ep_size` 参数来启用专家并行。大多数 MoE 模型默认采用 [token dispatch](#token-dispatch)，因此 `ep_size` 可以独立于 `tp_size` 进行设置。少数模型（如 Llama 4 和 Gemma 4）仍默认使用掩码和 all-reduce（即 `model.ep_plan` 中的 `"ep_router"` 和 `"moe_tp_experts"`）。任何具有 `ep_plan` 覆盖的模型也支持掩码，但要求 `ep_size=tp_size`，以确保专家组中的每个 rank 接收相同的 token。

```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 模型在其配置中定义了两个计划：用于稠密模块的 `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` 样式沿专家维度拆分专家权重，使每个 rank 仅加载其本地专家；`ep_router` 会在 all-reduce 合并专家输出之前，屏蔽位于其他 rank 上的专家。

`tp_plan` 仅在 `tp_size > 1` 时生效，`ep_plan` 仅在 `ep_size > 1` 时生效。当启用 TP 且禁用 EP 时，完整的 TP 计划将生效，包括专家规则。

`ep_plan` 中的专家前向规则决定了 token 如何到达各个专家：

| 规则 | 机制 | 布局 |
| :--- | :--- | :--- |
| 在路由器上使用 `"moe_tp_experts"` 和 `"ep_router"` | 掩码与 all-reduce：每个 rank 都在整个批次上运行其本地专家，路由器对其他专家进行掩码，然后通过 all-reduce 合并输出 | `ep_size=tp_size` |
| `"ep_dispatch_experts"` | [token dispatch](#token-dispatch)：每个 rank 保留自己的 token，仅通过两次 all-to-all 集合通信交换路由后的 (token, expert) 对 | `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/zh/expert_parallelism#transformers.DistributedConfig)，以覆盖预定义计划中的个别规则。未指定的规则将保留，合并后的计划会存储在模型上。每个键必须与某个模块、参数或现有计划条目相匹配；否则，在分片之前加载时就会引发 `ValueError`。请使用从已加载模型视角看到的完整路径，例如对于因果语言模型使用 `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`，并同时覆盖路由器和专家的前向规则（路由器模块名称因模型而异，例如在 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"` 以使用 token 分发。此时路由规则将被忽略，因为分发需要全局专家 ID 来查找每个专家的所属者。

## Token 分发

在 token 分发模式下，每个 rank 仅训练其负责的 batch 部分。在每个 MoE 层中，一个 rank 会对其 token 进行路由，通过 all-to-all 将每个 (token, expert) 对发送到拥有该 expert 的 rank，在接收到的数据上运行其本地 expert，再通过第二次 all-to-all 获取结果，并使用路由权重进行组合。只有被路由的激活值和 expert 输出会进行传输，且没有任何 rank 会计算其不拥有的 token 对应的 expert。

```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` 的限制。对于模型的其余部分：

- 专家之外的参数通过 [FSDP2](./fsdp) 在 `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/zh/main_classes/model#transformers.PreTrainedModel.save_pretrained) 会像处理其他参数一样将它们聚合。
- 一个专家并行组包含 `ep_size / tp_size` 个批次，因此某个专家的梯度是这些批次上的求和。`efsdp` 的归约操作除以的是 `fsdp_size` 而非其组大小，这与 FSDP2 对稠密模块所采用的每批次平均值相同。
- 每个本地专家在每一层中还会处理一行零填充。即使某个 rank 的专家未接收到任何 token，它仍会参与反向 all-to-all 通信以及专家梯度归约。

### 尺寸如何组合

- `tp_size * fsdp_size` 是进程数。`ep_size` 不会增加进程数：它仅针对专家权重重新组合相同的 rank。
- `ep_size` cuts the expert list into `ep_size` blocks. Each rank computes `num_experts / ep_size` experts, and `ep_size` consecutive ranks hold one complete set. That set of ranks is the group the all-to-all runs in.
- `efsdp_size = fsdp_size * tp_size / ep_size` 表示专家集的完整副本数量。位于不同副本中相同位置的 rank 会像 FSDP 对稠密模块所做的那样，对这些专家进行分片以节省内存，并对它们的梯度求平均。
- 某个 rank 持有的批次仅取决于 `fsdp`。连续的 rank 组成一个 TP 组并获取相同的批次；`fsdp_size` 个组则获取不同的批次。

八个进程，`tp_size=2, fsdp_size=4`，八个专家：

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

从图中可以得出两个数字：

- 在 EP 组内，每个 token 存在 `tp_size` 份，持有其批次的每一对 rank 各持有一份。这与 `ep_size` 无关。
- EP 组持有 `ep_size / tp_size` 个不同的批次。这是某个专家的梯度所累加的批次数量。

### 使用张量并行

将 `tp_size > 1` 设置为大于 1，以使用 TP 计划对稠密模块进行分片，同时让专家模块使用 dispatch。在八个进程上：

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

由于张量并行在 TP rank 对内部复制了激活值，因此每对 TP rank 接收相同的 batch。如果两个 rank 都派发了它们所有的 token，那么拥有某个专家的 rank 就会收到每个 token 两次，计算两次，其权重梯度也会翻倍。专家完整地存在于单个 rank 上，因此无法通过权重来拆分重复部分；而且拥有者通常是另一个 rank，因此也无法像 masking 那样通过所有权来解决。因此，该对 rank 会拆分行：每个 TP rank 派发互不重叠的 `1 / tp_size` 的 token，结果返回给发送它们的 rank，然后对零填充的半部分在该对 rank 之间执行 all-reduce，以恢复下一层所期望的复制输出。拆分依据是 `tp_size`，而不是 `ep_size`，因为只有持有 batch 的 rank 才能发送它。专家组跨越四个 rank，每个专家在 `efsdp_size = 2` 个 rank 之间进行 FSDP 分片，而主干的 FSDP 组跨越四个 rank。模型常规的 TP 约束（例如注意力头数的可整除性）仍然适用于稠密模块。Token 切片可能不均匀或为空，包括在单 token 解码期间。

Token dispatch 目前尚不能与流水线并行结合使用；请使用 `pp_size=1`（尚未测试）。这些配置均使用八块 GPU：

| 配置 | 结果 |
| :--- | :--- |
| `DistributedConfig(tp_size=4, fsdp_size=2, ep_size=4)` | 使用四个 TP 组进行分发，每个组将其批次切分为四份（默认计划）；除非你指定了 ep_plan 以使用旧版掩码 EP。 |
| `DistributedConfig(tp_size=1, fsdp_size=8, ep_size=4)` | 在每个 rank 上使用独立的 batch 进行分发，不进行切片，并将专家模型通过 FSDP 在 rank 对之间进行分片。 |
| `DistributedConfig(tp_size=2, fsdp_size=4, ep_size=4)` | 按每个批次一个 TP 对进行分发，每个对将其批次一分为二；每个专家组包含两个批次。 |
| `DistributedConfig(tp_size=8, ep_size=8)` | 通过让每个 rank 共享一个批次并将其切分为八份，或者通过掩码和 all-reduce 来实现掩码计划，来执行 Dispatch。 |

> [!WARNING]
> [Trainer](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer) 目前尚未考虑 token dispatch：batch 和 token 计数假设采用 all-reduce 布局，即 TP 组内的各个 rank 共享一个 batch，并且存在 `fsdp_size` 个数据并行分片。Trainer 对 dispatch 的支持将在后续版本中提供。

## 与 FSDP2 结合

Tensor 并行和专家并行将权重分片到 `tp` 上，但优化器状态以及没有规则匹配的模块仍然在组内的每个 rank 上复制，这限制了可训练模型的规模。在第二个 mesh 维度上添加 [FSDP2](./fsdp)，并设置 `fsdp_size`。通过掩码和 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/zh/main_classes/trainer#transformers.Trainer) 进行训练。它会计算两个 mesh 上的梯度范数，并为每个 mesh 分配独立的优化器参数组。[save_model()](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer.save_model) 会将分片权重聚合为一个常规检查点。这需要 `accelerate>=1.12`，以便 `Trainer` 能够将 `tp_size` 和 `fsdp_size` 镜像到 `~Accelerate.ParallelismConfig` 中。

下表对比了在 8xH100 GPU 上仅使用 EP 的训练与 2D EP+FSDP2 的训练。工作负载为在 bf16 精度、序列长度 2048 下对 Qwen3-30B-A3B 进行全量微调。更多的 FSDP 分片会降低峰值内存，但由于 FSDP2 会在 `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/zh/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, dp_replicate_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` 个 rank 上应用专家并行计划。显式指定的 `ep_size` 具有更高优先级。此标志不会更改 `tp_size` 或 `fsdp_size`。

fsdp_size (`int`, *可选*) : FSDP（数据并行）的设备数量。如果 `None` 且设置了 `tp_size`，则默认为 1。

dp_replicate_size (`int`, *optional*) : FSDP 副本数量（HSDP）。参数在 `fsdp_size` 个设备上进行分片，并复制 `dp_replicate_size` 次：梯度先在副本内部进行 reduce-scatter，然后在副本之间进行 all-reduce。默认值为 1。我们建议将其用于多节点训练：普通的 FSDP 会在每一步中跨所有设备对参数进行 all-gather 并对梯度进行 reduce-scatter，因此这些流量会经过较慢的节点间链路。将 `dp_replicate_size` 设置为节点数量，并将 `fsdp_size * tp_size` 设置为每个节点的设备数量：这样分片流量就会保留在快速的节点内链路上，只有梯度 all-reduce 会跨节点传输。代价是每个设备的内存占用会增加，因为参数仅在节点内进行分片。

fsdp_plan (`dict[str, str]`, *可选*) : FSDP2 分片计划。保持为 `None` 以使用模型预定义的 `base_model_fsdp_plan` 和类级别的 `_fsdp_plan`。传入字典以覆盖个别规则；未指定的规则将保留。仅在 `dp_replicate_size * fsdp_size > 1` 时应用。

fsdp_cpu_offload (`bool`，*可选*，默认为 `False`)：是否启用 FSDP2 的 CPU 卸载。

fsdp_mixed_precision (`bool`，*可选*，默认为 `False`)：是否启用 FSDP2 的混合精度。

pp_size (`int`, *可选*) : 用于流水线并行的设备数量。如果 `None` 且设置了其他并行模式，则默认为 1。

ep_size (`int`, *可选*) : 拥有不同专家分片的设备数量。默认值为 1。需显式设置以启用 EP。必须是 `tp_size` 的倍数，且能整除 `fsdp_size * tp_size`。All-reduce 专家方案要求 `ep_size=tp_size`；token 分发（`"ep_dispatch_experts"`）也允许 `ep_size > tp_size`。

ep_plan (`dict[str, str]`, *optional*) : Expert parallel sharding plan. Leave as `None` to use the model's predefined `base_model_ep_plan`. Pass a dictionary to override individual rules of that plan; unspecified rules are kept. Applied only when `ep_size > 1`, and its rules take precedence over `tp_plan` rules for the same modules. An `"ep_dispatch_experts"` rule selects all-to-all token dispatch instead of router masking and all-reduce.

用于原生分布式推理和训练的配置，支持张量、流水线或 FSDP2 并行。