虽然专家并行(EP)直接处理 MoE 层内专家的分布,数据并行(DP)复制整个模型,但流水线并行(PP)提供了一种补充方法来管理大型模型(包括包含 MoE 层的模型)可观的内存占用和计算图深度。流水线并行将模型的层(而非数据或专家本身)划分为顺序阶段,并将每个阶段分配给一组不同的处理设备。设想一个深度 Transformer 模型。与将所有层放置到每个设备上(如纯 DP)或仅将 MoE 专家拆分到不同设备上(如纯 EP)不同,PP 划分层序列。例如,第 1-8 层可能在设备组 A 上构成阶段 1,第 9-16 层在设备组 B 上构成阶段 2,依此类推。微批处理与流水线执行PP 的简单实现会通过阶段 1 处理完整数据批次,将激活数据传递给阶段 2,在那里处理,然后顺序继续。这会导致设备大量空闲时间,因为任何时刻只有一个阶段处于活跃状态。为提高利用率,PP 采用微批处理。完整数据批次被拆分为更小的微批次($m_1, m_2, ..., m_k$)。执行以交错方式进行:阶段 1 处理 $m_1$。阶段 1 处理 $m_2$,同时阶段 2 处理来自阶段 1 的 $m_1$ 输出。阶段 1 处理 $m_3$,阶段 2 处理 $m_2$ 的输出,阶段 3 处理 $m_1$ 的输出,等等。这产生了一个“流水线”效应,使多个阶段同时保持忙碌。前向传播将微批次传递通过各个阶段,反向传播则以相反顺序将梯度传回。digraph G { rankdir=LR; node [shape=box, style=filled, fontname="Arial", fontsize=10]; edge [fontname="Arial", fontsize=9]; subgraph cluster_fwd { label = "前向传播"; style=dashed; color="#adb5bd"; F1 [label="阶段 1\n(设备 A)", fillcolor="#a5d8ff"]; F2 [label="阶段 2\n(设备 B)", fillcolor="#96f2d7"]; F3 [label="阶段 3\n(设备 C)", fillcolor="#ffec99"]; F1 -> F2 [label="激活 (m1)"]; F2 -> F3 [label="激活 (m1)"]; } subgraph cluster_fwd_m2 { label = ""; style=invis; F1_m2 [label="阶段 1 (m2)", fillcolor="#a5d8ff", shape=plaintext]; F2_m2 [label="阶段 2 (m2)", fillcolor="#96f2d7", shape=plaintext]; F1_m2 -> F2_m2 [style=invis]; } subgraph cluster_fwd_m3 { label = ""; style=invis; F1_m3 [label="阶段 1 (m3)", fillcolor="#a5d8ff", shape=plaintext]; } subgraph cluster_bwd { label = "反向传播"; style=dashed; color="#adb5bd"; B3 [label="阶段 3\n(设备 C)", fillcolor="#ffec99"]; B2 [label="阶段 2\n(设备 B)", fillcolor="#96f2d7"]; B1 [label="阶段 1\n(设备 A)", fillcolor="#a5d8ff"]; B3 -> B2 [label="梯度 (m1)"]; B2 -> B1 [label="梯度 (m1)"]; } subgraph cluster_bwd_m2 { label = ""; style=invis; B2_m2 [label="阶段 2 (m2)", fillcolor="#96f2d7", shape=plaintext]; B3_m2 [label="阶段 3 (m2)", fillcolor="#ffec99", shape=plaintext]; B3_m2 -> B2_m2 [style=invis]; } subgraph cluster_bwd_m3 { label = ""; style=invis; B3_m3 [label="阶段 3 (m3)", fillcolor="#ffec99", shape=plaintext]; } # 用于对齐的不可见边(近似) F1 -> F1_m2 [style=invis]; F1_m2 -> F1_m3 [style=invis]; F2 -> F2_m2 [style=invis]; B1 -> B2_m2 [style=invis]; B2_m2 -> B3_m3 [style=invis]; B2 -> B3_m2 [style=invis]; # 时间线模拟节点 T0 [label="T0", shape=plaintext, group=t] T1 [label="T1", shape=plaintext, group=t] T2 [label="T2", shape=plaintext, group=t] T3 [label="T3", shape=plaintext, group=t] T4 [label="T4", shape=plaintext, group=t] T5 [label="T5", shape=plaintext, group=t] T6 [label="T6", shape=plaintext, group=t] T7 [label="T7", shape=plaintext, group=t] {rank=same; T0; T1; T2; T3; T4; T5; T6; T7;} # 连接时间线(使用不可见节点构建结构) P1_T1 [label="F1(m1)", fillcolor="#a5d8ff", group=t1]; P1_T2 [label="F1(m2)", fillcolor="#a5d8ff", group=t1]; P2_T2 [label="F2(m1)", fillcolor="#96f2d7", group=t2]; P1_T3 [label="F1(m3)", fillcolor="#a5d8ff", group=t1]; P2_T3 [label="F2(m2)", fillcolor="#96f2d7", group=t2]; P3_T3 [label="F3(m1)", fillcolor="#ffec99", group=t3]; P1_T4 [label="...", shape=plaintext]; P2_T4 [label="F2(m3)", fillcolor="#96f2d7", group=t2]; P3_T4 [label="F3(m2)", fillcolor="#ffec99", group=t3]; B3_T4 [label="B3(m1)", fillcolor="#ffec99", group=t3b]; P2_T5 [label="...", shape=plaintext]; P3_T5 [label="F3(m3)", fillcolor="#ffec99", group=t3]; B3_T5 [label="B3(m2)", fillcolor="#ffec99", group=t3b]; B2_T5 [label="B2(m1)", fillcolor="#96f2d7", group=t2b]; P3_T6 [label="...", shape=plaintext]; B3_T6 [label="B3(m3)", fillcolor="#ffec99", group=t3b]; B2_T6 [label="B2(m2)", fillcolor="#96f2d7", group=t2b]; B1_T6 [label="B1(m1)", fillcolor="#a5d8ff", group=t1b]; B2_T7 [label="B2(m3)", fillcolor="#96f2d7", group=t2b]; B1_T7 [label="B1(m2)", fillcolor="#a5d8ff", group=t1b]; B1_T8 [label="B1(m3)", fillcolor="#a5d8ff", group=t1b]; # 连接时间线事件 T1 -> P1_T1; T2 -> P1_T2; T2 -> P2_T2; T3 -> P1_T3; T3 -> P2_T3; T3 -> P3_T3; T4 -> P1_T4; T4 -> P2_T4; T4 -> P3_T4; T4 -> B3_T4; T5 -> P2_T5; T5 -> P3_T5; T5 -> B3_T5; T5 -> B2_T5; T6 -> P3_T6; T6 -> B3_T6; T6 -> B2_T6; T6 -> B1_T6; T7 -> B2_T7; T7 -> B1_T7; # 添加用于 T8 位置的不可见节点 T7_dummy [style=invis, group=t]; T7 -> T7_dummy [style=invis]; T7_dummy -> B1_T8; }包含 3 个阶段和微批处理的流水线并行可视化。F 表示前向传播,B 表示反向传播,m 表示微批次索引。空闲期(气泡)发生在调度开始和结束时。整合流水线并行与 MoE 层MoE 层如何与流水线阶段配合?通常,整个 MoE 层(门控网络和相关专家)位于单个流水线阶段内。由于将 token 路由到专家涉及复杂的依赖关系和通信,通常避免将 MoE 层本身跨越阶段边界进行拆分。因此,包含 MoE 层的流水线阶段内部结构可能如下:输入激活数据来自前一阶段(或初始嵌入层)。这些激活数据由当前阶段内部的层处理,可能包括标准 Transformer 块。如果存在 MoE 层:门控网络计算每个 token 的专家分配。All-to-All 通信在该阶段分配的设备组内发生,以将 token 路由到其指定专家(假设该阶段内也使用了 EP)。专家处理其分配的 token。输出被合并,并可能传递给同一阶段内的后续层。阶段的最终激活数据传递给下一个流水线阶段。混合并行策略:PP + DP + EP对于真正大型的 MoE 模型,流水线并行与数据并行和专家并行结合使用时最有效。常见配置是:流水线并行(PP): 模型的层被划分为跨不同节点或 GPU 组的阶段。这解决了激活内存瓶颈,并允许模型规模突破单节点内存容量的限制。数据并行(DP): 在每个流水线阶段内,分配给该阶段的层被复制到多个 GPU。每个 GPU 处理不同的微批次(或其中一部分)。这扩展了阶段内的计算。专家并行(EP): 对于流水线阶段内的任何 MoE 层,专家分布在该阶段的数据并行副本上。token 路由的 All-to-All 通信发生在该特定阶段中参与的 GPU 之间。digraph G { compound=true; node [shape=box, style=filled, fontname="Arial", fontsize=10]; edge [fontname="Arial", fontsize=9]; subgraph cluster_PP { label="流水线并行(跨节点/组)"; style=dashed; color="#868e96"; rankdir=LR; subgraph cluster_Stage1 { label="阶段 1"; bgcolor="#e9ecef"; node [fillcolor="#a5d8ff"]; S1_DP0 [label="GPU 0\n(DP 秩 0)\n层 1-8\n专家 0-N/2"]; S1_DP1 [label="GPU 1\n(DP 秩 1)\n层 1-8\n专家 N/2-N"]; S1_DP0 -> S1_DP1 [label="All-to-All (EP)", style=dashed, dir=both, constraint=false, color="#7048e8"]; } subgraph cluster_Stage2 { label="阶段 2"; bgcolor="#e9ecef"; node [fillcolor="#96f2d7"]; S2_DP0 [label="GPU 2\n(DP 秩 0)\n层 9-16\n专家 0-N/2"]; S2_DP1 [label="GPU 3\n(DP 秩 1)\n层 9-16\n专家 N/2-N"]; S2_DP0 -> S2_DP1 [label="All-to-All (EP)", style=dashed, dir=both, constraint=false, color="#7048e8"]; } # 流水线通信链接 S1_DP0 -> S2_DP0 [label="激活\n(PP 链接)", lhead=cluster_Stage2, color="#f76707", style=solid, minlen=2]; S1_DP1 -> S2_DP1 [label="激活\n(PP 链接)", lhead=cluster_Stage2, color="#f76707", style=solid, minlen=2]; # 用于 DP 对齐的不可见边 # S1_DP0 -> S1_DP1 [style=invis]; S2_DP0 -> S2_DP1 [style=invis]; } # 添加解释缩写词的说明 Note [shape=plaintext, label="DP: 数据并行\nEP: 专家并行\nPP: 流水线并行\n箭头表示主要数据流(激活/梯度)或 All-to-All。", fontsize=9]; }结合了 PP、DP 和 EP 的混合并行。阶段通过 PP 形成。在每个阶段内,DP 复制阶段逻辑,EP 将专家分布在 DP 秩上。EP 的 All-to-All 发生在一个阶段内部,而 PP 的激活/梯度则在阶段之间流动。MoE 模型中 PP 的注意事项流水线气泡: 尽管微批处理有帮助,但流水线启动和结束阶段仍然会引入空闲时间(气泡)。高级调度技术,如交错流水线调度(例如 1F1B 调度),可以通过更有效地重叠前向和后向传播来进一步缓解此问题。阶段负载均衡: 实现最佳吞吐量需要平衡流水线阶段之间的计算负载。MoE 层通常比标准密集层计算密集得多。放置 MoE 层需要仔细考虑,避免使某个阶段成为主要瓶颈。专家数量和专家大小影响阶段持续时间。内存均衡: 类似地,内存需求(参数和激活)应在阶段之间均衡。PP 在分配激活内存方面表现优异,激活内存会随着序列长度和批次大小扩展。MoE 层增加参数内存,在阶段内通过 EP 分配。通信: PP 引入了阶段间通信开销,用于激活(前向)和梯度(反向)。这通常是相邻阶段设备组之间的点对点或集体通信。必须权衡此通信成本与包含 MoE 层的阶段内部的 All-to-All 成本。优化两者非常重要。框架复杂性: 实现混合并行策略需要复杂的分布式训练框架,如 DeepSpeed、Megatron-LM 或 Tutel,这些框架提供抽象来管理 PP、DP 和 EP/TP(张量并行)之间复杂的交互。流水线并行,特别是当它整合到混合策略中时,提供了一种强有力的机制,使得 MoE 模型能够突破单节点内存和计算的限制,通过有效地在层和设备间划分工作负载,使得训练具有万亿参数的模型成为可能。