标准数据并行通过在每个设备上处理不同数据批次,可以有效地扩展多个设备的训练,但这要求模型的完整副本驻留在每个设备上。对于拥有数千亿甚至数万亿参数的专家混合模型来说,这是不可行的;完整参数集的内存需求远远超出任何单个加速器的容量。对于Mixture of Experts模型,由于其参数量非常大,标准数据并行训练方法会遇到内存限制,因为它要求每个设备存储模型的完整副本,这往往超出单个加速器的容量。专家并行是一种专门的分布式训练技术,旨在应对这类内存难题。专家并行不是复制整个模型,而是将专家分片到可用设备上。每个设备只拥有一部分总专家,而模型的非MoE部分(如自注意力块和门控网络)仍保持复制。All-to-All 通信模式专家并行的一个显著特点是它依赖于一种名为All-to-All的通信集合。当一个令牌需要由位于不同设备上的专家处理时,它必须通过网络进行物理发送。这在每个MoE层内分两阶段进行。令牌分派(第一次 All-to-All): 每个设备上复制的门控网络为其本地批次令牌确定专家分配后,会发生一次全局数据洗牌。每个设备将其令牌发送到拥有其指定专家的特定设备。结果聚合(第二次 All-to-All): 远程专家处理完其指定令牌后,生成的输出表示会发送回其原始源设备。此过程确保每个令牌都由正确的专家处理,无论该专家位于何处,并且令牌表示会返回其原始序列顺序,以继续前向传播。下图显示了数据在两个设备(GPU)间的流动,每个设备拥有一组不同的专家。digraph G { rankdir=TB; splines=true; node [shape=plaintext, fontname="Arial", fontsize=11]; edge [fontname="Arial", fontsize=10]; subgraph cluster_0 { label="设备 0"; style="rounded"; bgcolor="#f8f9fa"; GPU0 [label=< <TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" BGCOLOR="#e9ecef" STYLE="ROUNDED"> <TR><TD PORT="in0">输入令牌(本地批次)</TD></TR> <TR><TD>门控网络</TD></TR> <TR><TD>自注意力</TD></TR> <TR><TD PORT="gate0">路由逻辑</TD></TR> <TR><TD PORT="exp0">专家 0-7</TD></TR> <TR><TD PORT="out0">输出</TD></TR> </TABLE> >]; } subgraph cluster_1 { label="设备 1"; style="rounded"; bgcolor="#f8f9fa"; GPU1 [label=< <TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" BGCOLOR="#e9ecef" STYLE="ROUNDED"> <TR><TD PORT="in1">输入令牌(本地批次)</TD></TR> <TR><TD>门控网络</TD></TR> <TR><TD>自注意力</TD></TR> <TR><TD PORT="gate1">路由逻辑</TD></TR> <TR><TD PORT="exp1">专家 8-15</TD></TR> <TR><TD PORT="out1">输出</TD></TR> </TABLE> >]; } // 仅自处理循环 edge [style=dashed, color="#495057", penwidth=1.0]; GPU0:gate0 -> GPU0:exp0 [label="用于专家 0-7 的令牌"]; GPU1:gate1 -> GPU1:exp1 [label="用于专家 8-15 的令牌"]; } 专家并行中的数据流。本地令牌通过第一次All-to-All通信步骤(红色)路由到其他设备上的专家。处理后,结果通过第二次All-to-All步骤(蓝色)返回到其原始设备。性能影响和权衡专家并行巧妙地解决了内存问题,使得模型能够扩展到数万亿参数。然而,这项优势伴随着显著的通信开销。All-to-All操作是分布式计算中带宽消耗最大的集合操作之一,因为每个设备必须同时与其他所有设备通信。此操作的性能受限于设备间的互连带宽(例如,节点内GPU的NVLink,或多节点设置的网络结构)。洗牌的总数据量与令牌数量、它们的隐藏维度以及每个令牌被路由到的专家数量(Top-k路由中的 $k$)成比例。这带来一个重要的权衡:内存节省: 可以存储更大的模型。通信成本: 训练吞吐量现在对网络带宽很敏感,而不仅仅是对计算(FLOPs)能力敏感。在实践中,这意味着简单地增加更多设备并不能保证线性加速。All-to-All洗牌的效率成为总训练时间中的一个主要影响因素。代码实现虽然底层逻辑复杂,但为大规模训练设计的现代深度学习框架抽象掉了大部分难度。DeepSpeed或fairscale等库提供了自动处理专家分片和通信的模块。下面是分布式MoE层内部发生情况的简化伪代码表示。# 单设备上专家并行前向传播的伪代码 # 非MoE层照常处理 x = self_attention(input_tokens) # 门控网络在所有设备上复制 # 它计算其本地令牌的分配 gate_logits = gating_network(x) local_assignments = top_k_router(gate_logits) # 第一次 All-to-All:将令牌洗牌到拥有目标专家的设备。 # 框架处理底层通信。 # 每个设备接收一批旨在供其本地专家处理的令牌。 shuffled_tokens = all_to_all_dispatch(x, local_assignments) # 使用本地专家子集处理接收到的令牌。 # 这些专家只存在于此设备上。 processed_shuffled_tokens = local_experts(shuffled_tokens) # 第二次 All-to-All:将处理过的令牌洗牌回其原始设备。 # 这恢复了原始令牌顺序。 output_tokens = all_to_all_combine(processed_shuffled_tokens, local_assignments) # 模型的其余部分继续进行 ...这种模型分片策略是训练真正大规模MoE模型的根本。然而,它很少单独使用。为获得最高效率,专家并行几乎总是与其他并行技术结合使用,我们将在下一节讨论该主题。