专家并行：在不同设备上分配专家

扩展MoE模型会遇到很大的内存需求，这主要是因为单个专家网络的数量可能非常多。标准数据并行（DP）中，每个工作节点都保存模型的完整副本，但当专家数量增加到几十甚至几百个时，这种做法很快变得不切实际。每个专家本身可能是一个拥有数百万参数 (parameter)的多层感知器（MLP）。在每个设备上存储所有专家会导致过高的内存消耗。

专家并行（EP）直接解决了这个问题，通过将MoE层中的专家划分到可用的计算设备（例如GPU）上。每个设备不再拥有所有$N$个专家，而是只持有一部分，通常是$N/D$，其中$D$是参与专家并行组的设备数量。

专家并行工作机制

考虑Transformer块中的MoE层。在专家并行下，处理一批令牌的操作顺序包含以下步骤：

1. 本地门控： 每个设备通过Transformer块的共享组件，直到MoE层的门控网络，处理其本地部分的输入令牌。门控网络通常相对较小，一般会在所有设备上复制（类似于数据并行）。它计算其本地令牌的专家分配。
2. 令牌路由（全对全发送）： 这是EP中具有代表性的通信步骤。根据门控决定，每个设备确定其本地令牌中哪些需要由其他设备上的专家处理。使用All-to-All通信原语在设备间交换令牌。设备$i$将其分配给设备$k$上专家的令牌直接发送给设备$k$。
3. 并行专家计算： 令牌到达正确的设备后，每个设备使用其本地的专家子集计算所接收令牌的输出。由于每个设备只持有$N/D$个专家，内存占用量显著减少。这种计算在EP组的所有设备上并行进行。
4. 结果收集（全对全接收）： 专家计算完成后，处理过的令牌表示必须返回其原始设备，以保持后续层的序列完整性。另一个All-to-All通信步骤会收集这些结果，确保每个设备收到与它在步骤1中最初处理的令牌相对应的输出。

下图显示了四个设备上的流程，每个设备拥有两个不同的专家。

8个专家分布在4个设备上（每个设备2个专家）。虚线表示第一次All-to-All通信（根据门控分配将令牌$T$发送给目标专家$E_i$）。点线表示第二次All-to-All（将已处理令牌$P(T)$返回到它们的原始设备）。为清楚起见，仅显示部分路径。

优点与代价

专家并行主要的优点是内存效率高。通过划分专家，可以实例化总参数 (parameter)量远超单个设备容量的MoE模型。这使得在不按比例增加单个工作节点内存负担的情况下扩展模型容量（通过更多专家）成为可能。它还分配了专家前向和反向传播 (backpropagation)的计算负载。

然而，这种好处是以增加通信开销为代价的。两次All-to-All操作是通信密集型的，尤其是在大规模情况下。它们的延迟和带宽要求可能成为重要的瓶颈，如果管理不当，可能会限制整体训练吞吐量 (throughput)。优化这种通信是扩展MoE模型时的主要侧重点，本章后面会讨论。

实现时需考虑的方面

• 专家放置： 虽然将专家简单地以轮询方式分配给设备很常见，但更完善的策略可能会考虑网络拓扑结构，以尽量减少频繁通信设备之间的通信开销。
• 框架集成： 实现高效的All-to-All通信和管理分布式状态需要专门的库。DeepSpeed（使用其MoE实现）和Tutel等框架抽象了大部分复杂性，提供了优化的通信核和与PyTorch等标准深度学习 (deep learning)框架的集成。这些库处理令牌的洗牌和专家计算的协调。
• 与其他并行方式的配合： 专家并行通常与数据并行结合使用。常见的设置是使用DP复制非MoE层，同时使用EP在同一组设备上划分MoE层的专家。我们将在下一节考察这种集成方式。

总之，专家并行是扩展MoE模型的一种基础技术。它在设备间划分专家，通过减少每个设备的内存需求来实现庞大的模型尺寸，但也带来了显著的All-to-All通信开销，需要仔细考虑和优化。