专家并行：混合专家模型的分布式训练方案

FreeGuideOnline 最新 2026-06-14

专家并行：解锁混合专家模型的分布式训练

随着大语言模型向万亿参数迈进，传统的稠密模型（每个输入都需要激活全部参数）在计算和内存上变得难以承受。混合专家模型（Mixture of Experts, MoE） 提供了一种高效的解决方案：它仅激活部分参数来处理每个输入，从而以更低的计算开销实现更大的模型容量。然而，MoE 模型包含大量稀疏激活的“专家”子网络，这给分布式训练带来了全新的挑战。专家并行（Expert Parallelism） 正是为 MoE 量身定制的一种分布式策略，它让成百上千个专家跨 GPU 协同工作，实现近乎线性的扩展。本教程将从原理到实践，为你拆解专家并行的核心设计。

1. 混合专家模型为何需要专门的并行策略

要理解专家并行，我们先快速回顾 MoE 的基本结构和对训练的挑战。

1.1 MoE 的结构：稀疏激活与专家集合

一个标准的 MoE 层会替换 Transformer 中的前馈网络（FFN）。它包含三个核心组件：

  • 门控网络（Router/Gate）：一个轻量级的线性层，根据输入 token 的表示，计算每个 token 应该被发送到哪些专家。
  • 专家集合（Experts）：通常是多个独立的 FFN。例如，一个 MoE 层可能拥有 8 个、64 个甚至更多专家。
  • 稀疏激活：门控机制每次为每个 token 选择 Top-K 个专家（如 Top-2），只有被选中的专家才会计算，其余专家完全不被激活。这意味着每个 token 实际使用的参数量只是模型总参数的一小部分。

1.2 传统并行策略面临的困境

如果把一个拥有 128 个专家的 MoE 模型直接放到标准数据并行中训练，每个 GPU 都需要持有全部专家的副本。这会导致：

  • 显存爆炸：即使每个 token 只激活 2 个专家，所有 128 个专家的参数都必须驻留在每个 GPU 显存中，总参数量远超单卡容量。
  • 计算浪费：数据并行要求每个 GPU 完成全部专家的计算，但实际大量专家未被任何 token 选中，造成功耗浪费。

简单的模型并行（如张量并行）虽然能切割单个专家，但无法解决“所有专家都必须存在”的问题。我们需要一种新的并行维度——按专家维度切割模型，这就是专家并行的核心思想。

2. 专家并行的核心原理

专家并行将不同的专家子网络放置在不同的 GPU 上，每个 GPU 只负责存储和计算一部分专家。门控机制则充当调度中心，将 token 动态路由到其目标专家所在的设备。

2.1 基础流程：路由与 All-to-All 通信

在一次 MoE 层的前向传播中，专家并行的执行步骤如下：

  1. 门控决策：每个 GPU 上的 token 独立通过门控网络，得到匹配分数，并选出 Top-K 专家的索引。
  2. 构建分发方案：根据目标专家索引，每个 token 被标记为需要发送到哪个目标 GPU（持有对应专家的设备）。
  3. Token 重分布（All-to-All Dispatch）：通过一次全交换（All-to-All） 集体通信，各个 GPU 之间交换 token。原本位于 GPU-0 的 token 如果需要 GPU-2 上的专家计算，就会被发送到 GPU-2。
  4. 专家本地计算：重分布后，每个设备上只包含分配给其本地专家的 token。各设备独立对本地 token 执行专家 FFN 计算。
  5. Token 回传（All-to-All Combine）：计算完成后，通过第二次 All-to-All 通信，将 token 按照其原始位置送回。这保证了 token 序列的顺序得以恢复，继续后续的其它层计算。

整个过程的核心是两次 All-to-All 通信，这也是专家并行区别于其它并行策略的关键。

2.2 容量因子与 Token 丢弃

实际训练时，各个专家收到的 token 数量往往不均衡。为防止某专家被过量 token 淹没导致显存溢出，常引入容量因子（Capacity Factor）。其公式为：

专家容量 = round( (总 token 数 / 专家数) × 容量因子 )

如果发往某个专家的 token 数超过其容量，多余 token 会被直接丢弃（不参与该专家计算，或通过残差连接绕过 MoE 层）。容量因子通常设为 1.0~1.25，在负载均衡和计算精度之间取得平衡。

3. 专家并行与其它并行策略的混合编排

纯专家并行很少单独使用。在实际训练超大规模 MoE 模型时，它通常与数据并行、张量并行等组合，形成多维混合并行。

3.1 专家并行 + 数据并行：复合扩展

最常见的组合是将专家并行维度和数据并行维度叠加。我们可以这样划分设备网格：

  • 专家并行组：例如，4 个 GPU 组成一个专家并行组，共同持有全部专家，每个 GPU 持有 1/4 的专家。
  • 数据并行组：复制多个这样的专家并行组，形成数据并行维度。

这种设计使得：

  • 在组内，token 通过 All-to-All 跨专家分发。
  • 在组间，处理完全不同的数据批次，并在反向传播后对同一专家参数进行梯度 All-Reduce。

3.2 张量并行 + 专家并行：突破单专家瓶颈

当单个专家本身参数量巨大（例如一个专家 FFN 的 hidden size 极大），单卡甚至无法装下一个专家时，就需要引入张量并行。张量并行的切分维度发生在单个专家内部，将矩阵乘法分割到多个 GPU。此时，一个专家的计算由一个小型 GPU 组共同完成，而专家并行则在更大的专家维度上跨组分布。

著名模型如 GShard、Switch Transformer 都采用了“张量并行 + 专家并行 + 数据并行”的三维混合策略，成功将模型参数扩展到万亿级别。

4. 深度学习框架中的专家并行实现

现代框架已经开始原生支持专家并行，极大降低了使用门槛。

4.1 DeepSpeed-MoE

微软的 DeepSpeed 库提供了专门针对 MoE 的优化包 DeepSpeed-MoE。你只需在配置中添加 MoE 相关参数，并设置专家并行维度：

# DeepSpeed 配置文件示例片段
"moe": {
    "num_experts": 64,
    "top_k": 2,
    "expert_parallel_size": 8  # 按 8 路专家并行切分专家
}

它内部会自动处理 All-to-All 通信调度和负载均衡损失（辅助损失），甚至支持将专家卸载到 CPU 内存。

4.2 Megatron-LM 的专家并行支持

NVIDIA Megatron-LM 通过 --expert-model-parallel-size 参数启用专家并行，并与原有的张量并行、流水线并行无缝集成。用户只需在定义 Transformer 层时用其 MoELayer 替换标准 FFN 层，框架就能自动完成跨设备的专家分布式计算。

4.3 PyTorch 分布式原生支持（torch.distributed）

从 PyTorch 2.0 开始，DTensor 和 DeviceMesh 为自定义专家并行提供了灵活的原语。你可以直接构建设备网格，利用 torch.distributed.all_to_all_single 等接口手动实现路由逻辑。这对于研究新型 MoE 路由算法（如 Hash 路由、强化学习路由）的团队非常有用。

5. 专家并行的负载均衡挑战与优化

专家并行虽然解决了显存和计算总量问题，但“动态路由”引入了负载不均这一核心挑战。

5.1 辅助损失（Auxiliary Loss）

最常用的负载均衡手段是在训练目标中加入一个辅助损失。它惩罚专家接收 token 数量与均匀分布之间的偏差：

L_aux = α * num_experts * Σ f_i * P_i

其中 f_i 是分配给专家 i 的 token 比例，P_i 是门控给该专家的平均概率。训练时梯度会驱使门控网络让每个专家被均匀选中。

5.2 专家并行中的局部与全局均衡

在专家并行场景下，均衡要求不单是“全局均匀”，还要考虑每个设备（GPU）上的计算负载均匀。因为一个设备可能持有多个专家（如 8 个专家分给 2 个 GPU，每卡 4 个），即使专家个体间有轻微不均，只要每卡上专家的总 token 数接近，All-to-All 后的计算仍然高效。因此，在设备级进行负载监控与动态容量调整是专家并行的优化方向。

5.3 通信与计算重叠

两次 All-to-All 通信是主要的延迟来源。优化手段包括：

  • 将专家计算中的矩阵乘法与通信重叠（需框架支持融合算子）。
  • 编码通信路由时，使用 NCCL 的专用 All-to-All 原语，避免成对的 send/recv 造成的同步开销。
  • 对 token 进行打包排序，确保通信时数据连续，提升带宽利用率。

6. 总结与最佳实践

专家并行是 MoE 模型从实验室走向生产的关键加速技术。它通过将专家切分到不同设备，并利用 All-to-All 通信完成 token 的动态调度，成功实现了“模型参数总量极大，但每设备激活参数极少”的理想。

对于正在尝试专家并行的团队，建议遵循以下路径：

  1. 从小规模开始：先在少量 GPU 上验证 Expert Parallel + Data Parallel 组合的正确性，使用框架已封装好的 MoE 接口。
  2. 监控负载分布：不仅要看辅助损失，更要实时观察每个设备接收的 token 数量，必要时调整容量因子或启用动态路由。
  3. 优先使用成熟方案：DeepSpeed-MoE 和 Megatron-LM 已经过了大规模验证，其通信与均衡优化能为你节省大量试错时间。
  4. 警惕 All-to-All 瓶颈：在网络带宽有限的集群上，可考虑通过拓扑感知的编排（将专家并行组放在更高带宽的节点内）来缓解通信压力。

专家并行技术仍在快速演进，随机路由、专家选择 token 等新范式不断涌现。掌握其核心思想，你就打开了高效训练万亿参数模型的大门。