大语言模型并行策略（七）：专家并行

专家并行是针对 Mixture of Experts（MoE）模型的一种并行策略。MoE 模型的前馈网络中包含多个专家网络，每个专家网络都是一个小的前馈网络。在训练和推理过程中，输入数据会被路由到其中的一个或多个专家网络进行计算，最后将多个专家网络的结果汇总得到最终的输出结果。由于每个专家网络的计算是相互独立的，因此可以将不同的专家网络分布到不同的 GPU 上进行计算，从而实现专家并行。

Mixture of Experts (MoE)

MoE (Mixture of Experts) 是对 Transformer 模型的一种扩展，它将 Transformer 模型中的前馈网络（Feed-Forward Network, FFN）替换为多个专家网络（Expert Network），每个专家网络都是一个小的前馈网络。

图片来源：GShard

每次前向计算时，首先会通过一个路由网络（Routing Network）来选择最适合的专家网络，然后将输入数据路由到这些专家网络进行计算。最后会根据每个专家网络的权重将它们的输出结果进行加权求和，得到最终的输出结果。

下面是一个简单的示意图，展示了 MoE 的工作原理：

图片来自：Mixture of Experts in Large Language Models

每个 token 首先会经过一个路由网络（Router），路由网络就是一个 FFN 加上一个 softmax 层，输入为 token 的特征向量，输出一个概率分布，表示该 token 被路由到每个专家网络的概率，也表示每个专家对该 token 的权重。然后选择权重最大的 $k$ 个专家网络，将 token 的特征向量路由到这些专家网络进行计算，最后将这些专家网络的输出结果进行加权求和，得到最终的输出结果。这里每一个专家网络都是一个参数量较小的 MLP 模块。

图片来源：MEGABLOCKS

直觉上，MoE 模型就像是拥有一个专家团队，每个专家都有自己擅长处理的模式，而路由网络就是一个调度员，它会根据任务的特征选择最适合的专家来完成任务。非 MoE 模型中，前馈网络是一个单一的 MLP 模块。此时要想提高模型的表达能力，就需要增加 MLP 的参数量，这就像是让一个人去学习所有的知识，这难度很大。而 MoE 使用多个专家，每个专家网络只需要学习特定的模式，这一方面降低了学习的难度，另一方面也提高了模型的表达能力。

MoE 模型的一个重要特点是稀疏性（Sparsity），即每个 token 只会被路由到其中的一个或多个专家网络进行计算，而不是所有的专家网络都参与计算。这种稀疏性使得 MoE 模型在参数量大幅增加的同时，计算量不会增加太多。

MoE 模型的参数量分析

我整理了几个开源的基于 MoE 架构的模型的参数量和激活参数量的统计数据：

表格中概念的说明：

• 总参数：模型的总参数量，包括所有的路由专家、共享专家、注意力层、嵌入层等所有参数。
• 激活参数：单个 token 前向计算实际使用的参数规模。
• 专家数：路由专家和共享专家的总和，为每个 MoE 层中专家数量，不是整模型把各层相加后的数量。
• 单专家参数量：仅计算一个专家的参数量，主要是 MLP 中几个线性层的参数量。对 SwiGLU/GLU 专家，其参数量为 $3 \times d_{\text{model}} \times d_{\text{moe}}$，其中 $d_{\text{model}}$ 是模型的隐藏维度，$d_{\text{moe}}$ 是专家的隐藏维度。
• 总层数/MoE 层：模型中总的层数和 MoE 层数，有的模型在前几层使用 dense FFN，而后续层使用 MoE FFN。
• 专家参数占比：所有 MoE 层中全部路由与共享专家的 FFN 参数 / 官方总参数。

MoE 模型的总结

MoE 的总参数大多是专家参数，专家参数通常占总参数的 95% ~ 98%，而 Dense 模型中，FFN 的参数量通常占总参数的 50% ~ 90%。虽然 MoE 模型的总参数量很大，但每个 token 的前向计算只会使用其中的一部分专家参数，因此 MoE 模型的激活参数量通常比总参数量小得多。例如 DeepSeek-V3 为 671B 总参数、37B 激活参数。

最近几年的一个趋势是 MoE 模型的专家数量越来越多，单个专家的参数量越来越小。这种设计通常称为细粒度 MoE。在保持每个 token 激活的专家参数量大致不变的前提下，把一个大专家拆成多个小专家，可以显著增加可选专家组合的数量。例如，从 8 选 2 扩展到 128 选 8，路由器可以按 token 更细致地组合不同专家，使专家更容易学习特定领域、语言或模式，从而提升模型容量和专家专门化程度。

更小的专家还有助于提高负载均衡的灵活性：当某个专家过载时，路由器可以将部分 token 分配给功能相近的其他专家，而不必让一个很大的专家承担所有相关计算。细粒度专家也更适合分布到大量 GPU 上，使每张 GPU 可以承载多个独立的专家单元。不过，专家数增加并非没有代价：路由决策更复杂，跨 GPU 的 All-to-All 通信、负载均衡和小矩阵计算效率都会成为更突出的挑战。因此，专家数量、单专家大小和每个 token 的 Top-$k$ 通常需要结合模型质量、通信带宽、batch 大小与部署规模一起权衡。

专家并行的原理

前面介绍了混合专家（Mixture of Experts, MoE）模型的原理，不难发现每个专家实际上就是一个前馈网络（Feed-Forward Network, FFN），且每个专家的计算是相互独立的。单个 GPU 的显存容量可能无法容纳所有专家网络的参数，这时候就可以将不同的专家网络分布到不同的 GPU 上进行计算，每个 GPU 只负责计算自己负责的专家网络。

专家并行的原理很容易理解，就是将不同的专家网络分布到不同的 GPU 上进行计算，每个 GPU 只负责计算自己负责的专家网络。这个思路和张量并行类似，都是将模型的参数分布到不同的 GPU 上进行计算。不同的是，张量并行是将一个前馈网络的参数分布到不同的 GPU 上进行计算，而专家并行是将多个前馈网络分布到不同的 GPU 上进行计算。

专家并行通常会和张量并行结合使用，比如 TP=8 时，将 Attention 的参数分布到 8 个 GPU 上构成张量并行，同时将所有专家网络分布到 8 个 GPU 上构成专家并行。而 MoE 中的路由网络需要在每一个 GPU 上进行计算，因此路由网络的参数需要被复制到每一个 GPU 上。

通信模式

序列中的每个 token 需要根据门控机制选择最适合的专家网络进行计算。由于不同专家分布在不同的 GPU 上，每个 token 会被路由到不同 GPU 上的专家进行计算，因此几乎每个 GPU 都需要和所有其他 GPU 进行通信，才能将当前 batch 中的 token 发送到对应的专家网络。为了实现这一点，通常需要在 GPU 之间进行 All-to-All 通信，将当前 batch 中的 token 路由到对应的专家网络进行计算。

All-to-all 通信

All-to-all 可以让每个 rank 将自己的数据发送给所有其他 rank，同时接收所有其他 rank 的数据。每个 rank 都有一个发送缓冲区和接收缓冲区，发送缓冲区中存放着当前 rank 要发送给其他 rank 的数据，接收缓冲区中会存放其他 rank 发送给当前 rank 的数据。

下面是 all-to-all 通信的示意图：

以 rank 0 为例，起初它持有 1、2、3、4 这四个数据，分别需要发送给 rank 0、rank 1、rank 2、rank 3。当 all-to-all 通信完成后，rank 0 会收到 rank 0、rank 1、rank 2、rank 3 分别发送过来的数据，分别是 1、10、100、1000。

MoE 中的 all-to-all 通信

输入 token 首先会经过路由网络决定每个 token 被路由到哪个专家网络进行计算。由于每个专家网络需要处理的 token 数通常是不同的，因此每个 rank 需要发送和接收的 token 数也是不同的。

在执行 all-to-all 通信之前，需要预先分配好每个 rank 的发送缓冲区和接收缓冲区的大小。这里缓冲区的大小并不是以当前 rank 中待发送的 token 数量为准，而是以所有 rank 中待发送的 token 数量的最大值为准。因此首先需要进行一次 all-to-all 通信，相互之间交换每个 rank 中待发送的 token 数量，以此得到所有 rank 中待发送的 token 数量的最大值。

知道待发送的 token 数量的最大值后，就可以为每个 rank 分配发送缓冲区和接收缓冲区的大小。每个 rank 的发送缓冲区大小为 $ranks \times tokens \times dim$，其中 $ranks$ 是 rank 的数量，$tokens$ 是所有 rank 中待发送的 token 数量的最大值，$dim$ 是 token 的特征维度。

发送缓冲区分配好之后，就可以将当前 rank 中待发送的 token 数据拷贝到发送缓冲区中对应的位置，下面是一个示意图，展示了 rank 0 的发送缓冲区的内容：

然后执行 all-to-all 通信，将当前 rank 中待发送的 token 数据发送给其他 rank，同时接收其他 rank 发送过来的 token 数据。通信完成后，每个 rank 的接收缓冲区就存放了其他 rank 发送过来的 token 数据。之后就可以使用当前 rank 的专家网络对接收到的 token 数据进行计算，计算完成后再将结果发送回原来的 rank。

专家并行面临的挑战

专家并行的原理不难理解，但在实际应用中，专家并行面临着一些挑战。

挑战一：通信存在瓶颈

每次前向和反向计算都需要通过 all-to-all 通信来传递激活值或梯度，通信的延迟和带宽都会成为瓶颈，限制了专家并行的性能。另外，在选择专家时，需要依赖路由网络的输出结果，这导致计算与通信相互依赖，计算和通信无法重叠。

挑战二：负载不均衡

自然语言分布不均匀，高频词、通用语义会持续涌向少数几个专家，形成热专家，这导致热专家所在的 GPU 负载很高，其余 GPU 则很空闲，整体吞吐被最慢的 GPU 限制。

挑战三：计算不高效

因为一个专家网络实际上就是一个 MLP，如果 batch 数量较小，某个 GPU 上的每个专家网络可能只处理少量的 token，这会导致计算不高效。因为要完成 MLP 计算，需要读取专家网络的参数，但因为参与计算的 token 数量较少，无论采用什么优化方法，计算强度（arithmetic intensity）都不会很高，导致计算资源无法充分利用。

考虑只有一个 token 经过一个专家网络进行计算的情况：

对于 A100 GPU 来说，计算强度需要达到 100 以上，才能充分利用 GPU 的计算资源。

总结

专家并行是针对 MoE 模型的一种并行策略，它将不同的专家网络分布到不同的 GPU 上进行计算，从而利用多个 GPU 的计算资源来加速 MoE 模型的训练和推理。专家并行的原理比较简单，但在实际应用中面临着通信瓶颈、负载不均衡和计算不高效等挑战。