专家并行

专家并行是一种针对混合专家（MoE）模型的并行策略。每个专家的 feedforward 层位于不同的硬件加速器上。路由器将 token 分派给相应的专家并收集结果。由于每个 token 仅激活少数几个专家，该方法可以在不增加计算成本的情况下，将模型扩展到远更大的参数量。

DistributedConfig

使用 DistributedConfig 类和 ep_size 参数来启用专家并行。大多数 MoE 模型默认采用 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。

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 时，两者同时生效：张量并行 计划对注意力和稠密 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 如何到达各个专家：

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。请使用从已加载模型视角看到的完整路径，例如对于因果语言模型使用 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，并同时覆盖路由器和专家的前向规则（路由器模块名称因模型而异，例如在 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" 来使用 token 分发。此时，路由规则将被忽略，因为分发需要全局专家 ID 来查找每个专家的所属者。

Token 分发

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

• 专家之外的参数通过 FSDP2 在 fsdp 维度上进行分片，从而减少其梯度。
• 专家参数在 ep 上进行分片，并且当 efsdp_size = fsdp_size * tp_size // ep_size 大于 1 时，还会额外在 efsdp 上进行 FSDP 分片。它们始终被 FSDP 包装，因此 fsdp_mixed_precision 和 fsdp_cpu_offload 同样适用于它们，并且 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，八个专家：

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。在八个进程上：

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：

Trainer 目前尚未考虑 token dispatch：batch 和 token 计数假设采用 all-reduce 布局，即 TP 组内的各个 rank 共享一个 batch，并且存在 fsdp_size 个数据并行分片。Trainer 对 dispatch 的支持将在后续版本中提供。

与 FSDP2 结合

Tensor 并行和专家并行将权重分片到 tp 上，但优化器状态以及没有规则匹配的模块仍然在组内的每个 rank 上复制，这限制了可训练模型的规模。在第二个 mesh 维度上添加 FSDP2，并设置 fsdp_size。通过掩码和 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 进行训练。它会计算两个 mesh 上的梯度范数，并为每个 mesh 分配独立的优化器参数组。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 会略有下降。

对于在加载时进行分片的模型，尚不支持从检查点恢复，因此 Trainer 仅接受 save_only_model=True 或 save_strategy="no"。

API 参考