ラブパラ

\[\newcommand{\t}{\text} \newcommand{\bmat}[1]{\begin{bmatrix}#1\end{bmatrix}} \]

O. Introduction

本文计划系统性梳理各种并行方案以及它们的协同方式。

O.I. Notation

本章将所有 notation 统一。

并行相关配置：

  • \(N\)：总卡数。
  • \(N_\t{TP}\)：张量并行 (切分矩阵乘法) 数。
  • \(N_\t{PP}\)：流水线并行 (切分模型层结构) 数。
  • \(N_\t{CP/SP}\)：上下文并行 (切分输入序列长度) 数。需要注意，这里同时提到了 CP 和 SP，这两者有一些隐式的区别，本文虽然会遵循惯例区分二者，但读者如果不想区分也没有任何关系。
  • \(N_\t{DP}\)：数据并行 (切分 batch) 数。
  • \(N_\t{EP}\)：专家并行 (切分 expert) 数。
  • Dense 模型总有 \(N=N_\t{TP}\times N_\t{PP}\times N_\t{CP}\times N_\t{DP}\)，Sparse 模型还要额外乘一个 \(N_\t{EP}\)（有些 notation 中会将 DP 和 EP 会重叠，本文不这么做，具体将在下文展开分析）。

训练相关配置：

  • \(B\)：全局的等效 bsz。应在 optimizer 分析时使用。
  • \(b\)：单张卡上每次 BP 时的 bsz。应在单卡性能分析时使用。专业术语是 micro-batchsize。
  • \(m\)：gradient accumulate 步数。满足 \(B=b\times m\times N_\t{DP}\times N_\t{EP}\)。（因为 EP 总是伴随着 DP；有些 notation 会将二者混合为 \(N_\t{DP}\)，本文不这么做，具体将在下文展开分析）
  • \(S\)：文本长度。
  • \(T\)：单次 step 所经手的总 token 数。满足 \(T=S\times B\)。

模型架构相关配置：

  • \(L\)：模型层数。
  • \(H\)：attn/MoE/FFN 层的输入输出维度。
  • \(H_\t{FFN}\)：MoE/FFN 的内部维度。传统是 \(H_\t{FFN}=4H\)，如果用 SwiGLU 就是 \(\dfrac83H\)。
  • \(N_q\)：query 头数。
  • \(N_{kv}\)：kv 头数。\(N_q=N_{kv}\) 是标准 MHA，\(N_{kv}=1\) 是 MQA，\(1<N_{kv}<N_q\) 是 GQA。
  • \(d\)：单头维度。有 \(d=H/N_q\)。一般大模型 scale 宽度时都只增加头数，而 \(d\) 维持不变。
  • \(V\)：vocabulary size。

MoE 相关配置：

  • \(E\)：expert 数量。
  • \(K_\t{MoE}\)：每个 token 激活的 expert 数目。
  • \(E_\t{local}\)：每张卡上的 expert 数目。总是有 \(E_\t{local}=E/N_\t{EP}\)。

其它配置：

  • \(\delta\)：数据类型对应字节数。例如 FP32 时 \(\delta=4\)，BF16/FP16 时 \(\delta=2\)。
  • \(\Phi\)：总参数量。
  • \(\Phi_\t{act}\)：激活参数量。Dense 时 \(\Phi_\t{act}=\Phi\)，但 Sparse 时 \(\Phi_\t{act}\ll\Phi\)。
  • \(\beta_\t{inter}\)：跨节点带宽。
  • \(\beta_\t{intra}\)：节点内部带宽。
  • \(\pi\)：单卡峰值 FLOPs。

Optimizer 配置：

  • \(n_\t{opt}\)：optimizer 内部 tensor 数量。SGD 时 \(n_\t{opt}=0\)，动量法 \(n_\t{opt}=1\)，Adam/AdamW \(n_\t{opt}=2\)，Muon \(n_\t{opt}=1\)。
  • \(\delta_\t{opt}\)：optimizer 的字节数。即使模型本身量化了，为了保证稳定性，常常 optimizer 也必须有 \(\delta_\t{opt}=4\)。
  • \(P_\t{2D}\)：模型中 2D tensor 占比，因为 Muon 只对 2D 矩阵有效。
  • Adam 的总内存开销：
    • 如果 \(\delta=\delta_\t{opt}\)，此时最简单，每个参数需要 \(\delta\) 的存储、\(\delta\) 的梯度、\(2\delta_\t{opt}\) 的 Adam 一二阶动量，极限情况是 \(4\delta\Phi\) Bytes。
    • 如果 \(\delta\neq\delta_\t{opt}\)，则情况更为复杂：参数本身需要 \(\delta\) 存储 + \(\delta\) 梯度，此外 Adam 中还要存储一份精度为 \(\delta_\t{opt}\) 的副本，再加上原本的 \(2\delta_\t{opt}\) 一二阶动量（因此总 Optimizer State 是 \(3\delta_\t{opt}\)），极限情况是 \((2\delta+3\delta_\t{opt})\Phi\) Bytes。
  • Muon 同理，在能用 Muon 的地方只需要存一阶动量，不能用的地方就回滚到 Adam，如果量化就还要额外存高精度副本。此外，Newton-Schultz 一般还需要额外辅助变量，导致额外峰值开销。
  • 下文默认使用混合精度训练以及 \(P_\t{2D}\approx1\)，因此总是使用 \((2\delta+(n_\t{opt}+1)\delta_\t{opt})\Phi\) 作为 baseline。

O.II. Collective Communication

复习所有常见通信方式。认为每个节点在内存中持有一个 shape 相同的向量，它们构成一个通信组。

  • Reduce：所有节点的向量一起做某种按位运算（比如说加法/max），然后覆盖某特定节点的对应内存。
  • All-Reduce：所有节点一起 Reduce，用相同的聚合结果覆盖所有节点的对应内存。
  • Broadcast：某个特定节点用它的向量覆盖所有其它节点对应的内存。
  • All-Gather：将向量切分为 Group Size 块，然后第 \(i\) 个节点出第 \(i\) 块，把所有节点的块拼接起来覆盖所有节点的内存。
  • Reduce-Scatter：等效效果是先做 All-Reduce 然后第 \(i\) 个节点只保留第 \(i\) 块。实际有更高效的方法。事实上，All-Reduce 才是那个会被用 Reduce-Scatter + All-Gather 实现的 Primitive。

术语澄清：

  • 分片 (Sharded) 意味着只考虑该节点对应块的数据，其它数据不管。与之对应的是 全量 (Full/Replica) 数据。
  • 局部 (Local) 意味着这个数据未经过 reduce。而 全局 (Global) 意味着数据经过 reduce 了。

因此：

  • Reduce：把所有节点的局部全量信息集中到某特定节点，得到该节点上的全局全量信息。
  • Broadcast：把某特定节点的数据分发给所有节点。
  • All-Reduce：所有节点，局部全量 -> 全局全量。
  • All-Gather：所有节点，分片 -> 全量。
  • Reduce-Scatter：局部全量 -> 全局分片。

在并行性能分析中，我们一般忽略拓扑结构，使用单卡总通信量作为 metric。比如 All-Gather，虽然此时每个节点仅贡献 \(\Phi/N\) 的输出，但如果使用 Ring 等方法，实际总输出量其实是 \(\Phi\)（准确地说，\(\dfrac{N-1}N\Phi\)，不过不妨忽略这个系数）。或者在此时也可以使用输入量进行分析。

I. Parallelisms By Themselves

本节考虑每一种并行方法单独使用的效果。

I.O. Activations

突然发现自己对 autograd 框架的理解还是不足，这里补充一些不加并行时的细节。

裸的 pytorch autograd：

  • 在 Forward 时，计算出激活值 \(y=\t{op}_W(x)\) 后，会根据 \(\t{op}\)、参数 \(W\) 或输入 \(x\) 是否 requires_grad 来判断哪些值要保留。被保留的值会留在显存中。
  • 在 Backward 时，按照拓扑逆序遍历计算图，如果某个激活值的 BP 计算完成且已经完全传播给前一层，则立刻从显存中释放。
  • 显存峰值出现在 F 结束、B 开始的时刻；最终计算得到的梯度形状和参数形状相同，所有激活值及其梯度都会被释放。

FlashAttention：

  • 按照裸 autograd，\(S\times S\) 的 \(QK^\top\) 激活值是必须要保留的。
  • 因此在 F 时，一行行算，算完即弃，即可用 \(O(S)\) 的峰值显存算到最终输出 \(X\)。F 时也只往显存中写少量的统计量。
  • B 时重新计算 \(QK^\top\)，且仍然使用逐行计算模式。

Activation Checkpoint (AC)：

  • 把输入分层，仅记录每层入口处的激活值，中间的激活值算完直接释放。
  • BP 时利用入口处激活值重新计算中间激活值。
  • FlashAttn 利用了 AC 思想，但高度融合，更高效。
  • 代价是每一段内部都要重新跑 Forward，典型的 计算换存储。
  • 各种框架中都有手动/自动/选择性配置 AC 的方法。

I.I. Data Parallelism

回忆：

  • \(B\)：全局的等效 bsz。应在 optimizer 分析时使用。
  • \(b\)：单张卡上每次 BP 时的 bsz。应在单卡性能分析时使用。
  • \(m\)：gradient accumulate 步数。满足 \(B=b\times m\times N_\t{DP}\)。（本节默认 \(N_\t{EP}=1\)，也即只考虑 dense 场景）

因此：

  • 每张卡上都要存储完整的、大小为 \(\Phi\) 的模型，以及完整的若干 optimizer state。
  • 每次通讯，每张卡需要把本地梯度信息 all-reduce 式通讯。

所以有 ZeRO 方法，分三步：

  • ZeRO-0 就是上述裸 DP，一般通过 Reduce-Scatter + All-Gather 实现 All-Reduce。
  • ZeRO-1 把 optimizer state 分散到所有节点上，此时需要把梯度 reduce-scatter 到对应 optimizer state，再把更新结果 all-gather 回来。形式化地：
    • 每个节点上都有全量 \(\delta\Phi\) Bytes 参数。
    • 每个节点分别 BP，得到局部全量 \(\delta\Phi\) Bytes 梯度。
    • 梯度进行 Reduce-Scatter，局部全量->全局分片。
    • Optimizer State 被全局分片存储，每个节点只需存储 \((n_\t{opt}+1)\delta_\t{opt}\Phi/N_\t{DP}\) Bytes。
    • 全局分片梯度更新每个节点上存储的分片 Optimizer State 为最新状态，同时更新参数的对应分片，让分片参数达到最新状态。
    • 把参数 All-Gather 广播，全局分片->全局全量。
    • 它的单卡存储开销减为 \((2\delta+(n_\t{opt}+1)\delta_\t{opt}/N_\t{DP})\Phi\) Bytes。通讯开销 不变 (仍然是 Reduce-Scatter + All-Gather 各一次)。因此，在忽略拓扑、调度和 overlap 等工程细节的场合，ZeRO-1 是 免费 的。
    • 现有框架可以完美实现 ZeRO-1，手动配置是可选项。
  • ZeRO-2 把梯度也分散。
    • ZeRO-1 选择等局部全量梯度算完再 Reduce-Scatter。ZeRO-2 则不然，如果梯度中某一项 在所有卡中均计算完毕，则它可以立刻被 reduce 到对应卡上，其它卡则可以释放对应项的内存。
    • 特别地，需要区分 激活值的梯度 和 参数的梯度。前者 pytorch 会自行处理并释放，后者才是 DP 中真的需要 reduce 的东西。ZeRO-2 相当于在计算图中隐式加了一个「将所有卡的参数梯度 reduce 到对应卡」的步骤。
    • 进一步，为了避免 reduce 过于频繁，常见操作是把梯度打包为 bucket，等一个 bucket 中的所有梯度均 resolved 后再一起 reduce。bucket 如果没分好，也可能会导致等待某个梯度而迟迟不释放，拉高峰值内存占用。
    • 理想状态下，它可以重叠计算与通讯：上一层的 bucket 算完了正在 reduce，下一层的 bucket 正在计算。
    • 它的单卡存储开销是 \((\delta+\delta/N_\t{DP}+(n_\t{opt}+1)\delta_\t{opt}/N_\t{DP})\Phi\) Bytes，可能会有额外峰值开销，包括 activation 和 bucket 两部分。通讯开销本质上仍然是 Reduce-Scatter + All-Gather（前者被拆得零碎但总量不增加），因此也可以基本上看做 免费。
    • 现有框架通常已经实现了 ZeRO-2，但仍然需要手动配置 bucket size、grad accu 等参数。
  • ZeRO-3 (a.k.a. FSDP, Fully-Sharded DP) 把模型也分散。
    • 简单来说，Forward 时要从所有节点上 gather 需要的参数，Backward 时同理。
    • 单卡存储开销是 \((2\delta/N_\t{DP}+(n_\t{opt}+1)\delta_\t{opt}/N_\t{DP})\Phi\) + 额外峰值开销。通讯开销除了对梯度的 Reduce-Scatter 以外，每个节点还要额外在 Forward 和 BP 时各产生 \(\delta\Phi\) 的通讯，相当于 用通讯换存储。不过也有好处，就是省掉了最后模型参数的一次 All-Gather，因此总单卡通信量是 \(3\delta\Phi\)，是 ZeRO-2 的 \(1.5\) 倍。
    • gather 参数一般以层为单位，而层的划分需要 手动设定。具体地，只要指定划分的边界（比如说 TransformerBlock）即可。
    • 划分粒度越细，all-gather 的峰值显存开销就越小，但通讯频率上升，计算可能较难掩盖通讯。
    • 与此同时，跨层的调用会非常讨厌，甚至会让自动划分 policy 失效，因此尽量不要跨层访问 参数。（至于激活值？torch 自己会处理的！）
    • 可以启用 pre-fetching 以重叠通讯与计算。pre-fetching 量可以自行配置，相当于是某种程度的 存储换效率。
    • 在跨机柜的场合，ZeRO-3 需要太多通讯，效果并不好，因此可能会采取在机柜内部 ZeRO-3，但是跨机柜时回退到 ZeRO-2 甚至 1。
  • ZeRO-R 把激活值也分散。
    • 首先，裸的 ZeRO-3 有搭配或不搭配 AC 两种。前者会引入额外的计算开销，后者的存储开销较大。
    • ZeRO-R 搭配 AC 使用。但问题是，每个节点收到的 micro-batch 不同，激活值也不同，分散激活值不是脱裤子放屁吗？这样做第 \(i\) 个节点不能只存一份「所有节点共用」的第 \(i\) 段激活值，要分别存储每个节点的第 \(i\) 段激活值。
    • 事实上 ZeRO-R 的好处主要有两方面：一方面能让负载更均衡一点（尤其是搭配其它复杂并行方法的场合），另一方面虽然在 DP 时无法共用 AC，但其它并行中是可以共用的，尤其是 TP。

总结：

方法 并行对象 理论单卡存储 (Bytes) 单卡总收发量
0 / \((2\delta+(n_\t{opt}+1)\delta_\t{opt})\Phi\) \(2\delta\Phi\)
1 Optimizer State \((2\delta+(n_\t{opt}+1)\delta_\t{opt}/N_\t{DP})\Phi\) \(2\delta\Phi\)
2 Optimizer State
+ Gradient		 \((\delta+\delta/N_\t{DP}+(n_\t{opt}+1)\delta_\t{opt}/N_\t{DP})\Phi\) \(2\delta\Phi\)
3 Optimizer State
+ Gradient
+ Param		 \((2\delta/N_\t{DP}+(n_\t{opt}+1)\delta_\t{opt}/N_\t{DP})\Phi\) \(3\delta\Phi\)
R Activation / /

特别地，纯粹 DP 也即 ZeRO 1/2/3 都不管激活值；激活值占用的存储在 ZeRO-0 时或许无伤大雅，但可能 ZeRO-3 就会成为 bottleneck。因此有一些计算图比较奇葩的模型（比如最新最热的 AttnRes）的并行潜力会有限。一种可能的解决方案是继续分块，如果单个计算图比较怪，那就把整个计算图的深度减小，当成一个大 block，然后以传统方式堆叠多个这样的块。

特别地，grad accumulate 与 DP 有一些奇妙的互动：

  • ZeRO 0/1 时，本地有完整的梯度，因此可以正常 accumulate，它既能扩大全局等效 bsz \(B\)，还能降低中途 reduce-scatter 梯度的开销。实际流程是 \(m\) 次 BP -> 一次 Reduce-Scatter -> 一次 All-Gather。
  • 但是 ZeRO 2 时，本地没有完整的梯度，要想保留梯度切分的优势就只能每算一轮就 reduce 一次梯度。实际流程是 \(m\) 次 BP (每次内部都要 Reduce) -> 一次 All-Gather，仅仅省下了 All-Gather 的代价。或者也可以付出额外存储缓存梯度，相当于回退到 ZeRO-1。
  • ZeRO 3 时效果更差，无法节省任何通讯，唯一的效果是将等效 bsz 扩大 \(m\) 倍（但单个 batch 的计算开销也扩大同等倍数）。

特别地，ZeRO 应该被理解为一种对 必须在所有节点中同步的参数 的通用处理方式。在 CP 和 EP 中，因为也有相应需求，因此 ZeRO 也会被使用。

I.II. Pipeline Parallelism

数据并行时，ZeRO 0/1/2 均无法真正 scale 模型 size；ZeRO 3 最终激活值会成为 bottleneck（不管你有没有用 AC 或 ZeRO-R 等 trick）。

PP 选择把模型按照层数进行切分，每张卡持有模型的若干层。但是这样做就有问题——单次 BP 的计算流程如下：

Layer4       F1B1
Layer3     F1    B1
Layer2   F1        B1
Layer1 F1            B1

显然这样做并没有省计算量，反而引入了额外的传输开销。因此 PP 只有搭配 GA 才有用，此时才会出现真正的、典型的流水线重叠行为

Block4       F1F2F3F4B4B3B2B1       Update
Block3     F1F2F3F4    B4B3B2B1     Update
Block2   F1F2F3F4        B4B3B2B1   Update
Block1 F1F2F3F4  "Bubble"  B4B3B2B1 Update

按照前文的 notation，可以发现虽然并行度 \(N_\t{PP}=4\)，但其中有 \(\dfrac{N_\t{PP}-1}{m+N_\t{PP}-1}\) 的比例是 bubble。也就是说只有把 GA step \(m\) 开大，才能摊平 bubble 大小。

上述模式是经典的 GPipe 模式。这样做的优势在于：

  • 天然可以把每一层的参数 + 梯度 + 激活值 + 优化器状态全都分割，尤其是激活值比 DP 更能切。
  • DP 需要很多全局通讯（reduce-scatter, all-gather），但 PP 的通讯都是点对点的，对路由压力更小。正因如此，常见做法是把 PP 放到跨节点通讯中，节点内部高速通讯留给 TP。而且 PP 的通讯量级也更轻。

但有两个最大痛点：

  • \(\dfrac{N_\t{PP}-1}{m+N_\t{PP}-1}\) 的 bubble 太大。
  • 需要把 Forward 时所有 step 的激活值（不管 Block 内部有没有加 AC）都存储直到 BP，也就是说激活值的显存开销严格翻了 \(m\) 倍（DP + GA 因为是 F-B 交替所以不会引入额外显存开销）。

所以有改进后的 PP 方式，比如说 1F1B

Block4       F1B1F2B2F3B3F4B4F5B5F6B6
Block3     F1F2F3B1F4B2  B3F5B4F6B5F7
Block2   F1F2F3F4  B1  B2F5B3F6B4F7B5
Block1 F1F2F3F4      B1F5B2F6B3F7B4F8

1F1B 的核心思想是，一旦有 B 可以执行就立刻执行，因此才会看到 Block3 在 F4 之前就执行 B1。更精妙的一点是，如果 FB 耗时相等，则可以如上述示意图一样，除了开头结尾有少量 bubble 以外，后期没有任何 bubble。但是，如果把所有 bubble 加和，可以发现 1F1B 的总 bubble 数是不变的，优势仅在于对显存的节省。

Megatron 默认使用 1F1B，且支持更精妙的 Interleaved 1F1B，让 pipeline 级别数大于 \(N_\t{PP}\)，一张卡执行多个阶段。比如说以 8 stage 4 卡为例，则卡 1 会处理 S1 + S5，卡 2 处理 S2 + S6，以此类推。

Node4           F14F24F34F44F18B18F28B28F38B14B38B24F48B48B34
Node3        F13F23F33F43F17F27F37B17F47B27   B13B37B23   B47
Node2     F12F22F32F42F16F26F36F46   B16   B26   B12B36B22
Node1  F11F21F31F41F15F25F35F45         B15   B25   B11B35B21

它的 bubble 数明显更小，而且并没有打破拓扑的单点性，一个朴素的 ring topology 就能胜任。代价是需要更多传输。

还有一些 pipeline 切得更细，比如说区分对激活值的 BP 和对参数的 BP：前者在计算图中，必须优先处理才能让后面的任务继续；后者随便什么时候都可以处理，只要在参数更新前进行即可，因此可以随时拿出来填充流水线（当然，代价是其使用的激活值必须长期处于显存中，直到对应的参数更新处理完毕才能释放）。这些方法是 Zero Bubble PP 等框架的核心。

主流框架如 megatron 一般支持全自动的 PP 部署，只需要填几个参数（例如 PP 几层，用不用 interleaved）。但是这样做对于 hetero 的层（比如说首层 embedding 和末层 LM head）并不公平，且如果启用 weight tying（这两层共享权重）则会导致层切分逻辑更加复杂，需要更精细的逻辑。

I.III. Tensor Parallelism

PP 在深度维度切层，TP 在宽度维度切矩阵。具体地，一个典型的 TP 逻辑如下

graph subgraph "Y=GeLU(XA) → [Y_1,Y_2]=GeLU(X[A_1,A_2])" X[X] --> f["(identity,all-reduce)"] f --> X1[X] f --> X2[X] X1 --> XA1[XA1] X2 --> XA2[XA2] XA1 --"GeLU"--> Y1[Y1] XA2 --"GeLU"--> Y2[Y2] g["(all-gather,identity)"] Y1 --> g Y2 --> g g --> Y[Y] end

其中 (U,V) 意味着在 Forward 时当作 U，但 Backward 时当作 V。这个逻辑是 SPMD 的，并不需要 host 节点操纵。

同理也可以把 \(A\) 纵向切成 \(\bmat{A_1\\A_2}\)，有相似的逻辑。

进一步，如果连续乘多个，比如说 FFN 中经典的 \(Z=\t{GeLU}(XA)B\) 逻辑，则如果已经有切分 \(Y=\bmat{Y_1&Y_2},B=\bmat{B_1\\B_2}\)，则直接有 \(Z=Y_1B_1+Y_2B_2\)。因此经典的 FFN 层结构一般如下：

graph subgraph "Y=GeLU(XA) → [Y_1,Y_2]=GeLU(X[A_1,A_2])" X[X] --> f["(identity,all-reduce)"] f --> X1[X] f --> X2[X] X1 --> XA1[XA1] X2 --> XA2[XA2] XA1 --"GeLU"--> Y1[Y1] XA2 --"GeLU"--> Y2[Y2] end subgraph "Z=Dropout(YB) → Z=[Y_1,Y_2][B_1\\B_2]=Y_1B_1+Y_2B_2" Y1 --> Y1B1[Y1B1] Y2 --> Y2B2[Y2B2] Y1B1 --> g Y2B2 --> g g["(all-reduce,identity)"] g --Dropout--> Z[Z] end

一般 TP 只会选择在行或列切一次，而不会都切（不像 GPU 内部的 GeMM 会选择都切)。同时 TP 最多支持连续两次矩阵乘法（第一次竖着切，第二次横着切），超过两次的 TP 虽然也不是没有，但必须在中间插入额外通信层。

至于对 attn 的 TP，因为 attn 的 scale 是通过塞更多头进行的，头内部的 latent dim 保持不变，所以直接把不同头塞到不同卡上即可，一般总是塞得下的（塞不下也会选择结合 CP 而不是在 latent 方向拆）。一般也不会有人对头内部再划分，主要是因为划分后 softmax 就会强制要求 all-reduce 式通讯。如果使用 GQA 等，要确保所有同一组的 Q head 和 KV head 能塞到同一张卡上。因此，attn TP 的流程如下：

graph LR X["X:(B,S,H)"] --> f["(identity,all-reduce)"] f --> X1["X:(B,S,H)"] f --> X2["X:(B,S,H)"] X1 --GeMM--> QKV1["QKV1:(B,N_q/N_TP,S,d)"] X2 --GeMM--> QKV2["QKV2:(B,N_q/N_TP,S,d)"] QKV1 --"self-attn"--> Y1["Y1:(B,N_q/N_TP,S,d)"] QKV2 --"self-attn"--> Y2["Y2:(B,N_q/N_TP,S,d)"] Y1 --"reshape&proj"--> Z1["Z1:(B,S,H)"] Y2 --"reshape&proj"--> Z2["Z2:(B,S,H)"] Z1 --> g Z2 --> g g["(all-reduce,identity)"] g --Dropout--> Z[Z]

TP 对于每种新算子都要手动实现，不过一般会选择直接调 colwise TP 和 rowwise TP 两种 primitive 组合实现，不需要手搓内部细节。

TP 的特性：

  • 优势：没有 bubble，且不需要对训练流程的侵入性实现，只需要把某个算子覆盖。也不需要大 GA step，同时和 PP 一样能对激活值做很好的分摊。
  • 劣势：比 PP 的 layer-wise 通讯要高，且需要 all-reduce 而不是 p2p 通讯。

I.IV. Context/Sequence Parallelism

CP/SP 选择在 context 维度进行切分。

首先考虑 FFN/MoE 的 CP。因为 FFN 本身是 tokenwise 独立的行为，所以跑 CP 是简单的，只需要保证所有节点的 FFN 参数同步即可。而这正是 DP ZeRO 的功能。因此，FFN 的 CP 直接套 DP 即可。

然而，在 TP 中我们没有深入 attn 内部；但如果要做 CP，就需要相关考量。

因此引出 CP 的两种主流手法：

  • DeepSpeed Ulysses，其实就是在 FFN 阶段跑正常的 CP，attn 时把节点上持有的「所有 head 和 \(1/N_\t{CP}\) 的 context」数据变成「\(N_\t{head}/N_\t{CP}\) 的 head 和完整的 context」，然后在 attn 内部跑 TP；下一次 FFN 再次切换。形式化地：

    • 首先求 QKV。这仍然是一个 tokenwise 操作，而且需要保证同一个 CP 组内的 QKV proj 矩阵相同，因此仍然需要 ZeRO 保证同步。
    • 然后用 a2a 把 head 放到一起，内部跑 head 流程。
    • 跑完后再 a2a 恢复 context-wise 的分布，head 之间 concat。
    • 最后跑 proj，proj 仍然需要 ZeRO 同步。
    • 显而易见的，其需要大量的 a2a 通讯，且仍然需要保证单个卡塞得下整个 QKV group。在某些语境中，会用 SP 来特指 Ulysses 式方法，因为它并没有真正进入 head 内部。

  • Ring Attention，深入 head 内部。

    • 每张卡持有当前 context 在所有 head 上的 \(QKV\)。
    • 对于每个 head，使用类似 flash-attn 的 online softmax 手法，让 \(K,V\) ring 式地进行一圈遍历，即可得到 self-attn 结果。
    • 在使用 flash-attn 时，单卡峰值显存是 \(O(Sd/N_\t{CP})\)，而单卡计算复杂度是 \(O(S^2d/N_\t{CP})\)；单卡通信开销是 \(O(Sd)\)。
    • 因为使用 ring 通讯方法，所以只需要 ring topology 和 p2p 通讯，不像 Ulysses 需要 a2a。同时 ring 式结构也让其通讯容易被计算掩盖。一般认为 Ring Attn 是超长文本角度最有扩展性的并行方式，且有时用 CP 是特指 Ring Attn。

对比二者：

  • 通信角度：Ulysses 的 a2a 对同一时刻多对多传输的带宽要求很高，且通讯与计算是串行的，难以被掩盖，依赖高带宽通讯（但如果带宽高，效果会优于 Ring Attn）；Ring Attn p2p 路由更容易、带宽要求更低，且容易被掩盖。
  • 算子效率：Ulysses 是 head-wise 处理，没有破坏 head 内部结构，可以完美调用单卡中优化到极致的 flash-attn 算子；Ring Attn 在 \(S/N_\t{CP}\) 过小时可能会导致负载不满。
  • 显存瓶颈：Ulysses 有上限，而 Ring Attn 理论上没有（实际上仍然会受到其它工程参数的制约）。
  • 负载均衡：Ulysses 天然均衡，而 Ring Attn 尤其在有 causal mask 时不均衡，相关的调度优化会比较复杂。
  • 实现难度：显然 Ulysses 实现难度远低于 Ring Attn。

I.V. Expert Parallelism

[!NOTE]

MoE 的 notation 在不同场合区别较大。本文认为 \(N_\t{DP}\) 和 \(N_\t{EP}\) 是两个独立的维度，但 EP 总是伴随着相应的 DP，因此 \(B=b\times m\times N_\t{DP}\times N_\t{EP}\)。但需要注意，这仅仅是本文（以及其它一批资料）的 notation，其它材料可能使用不同的 notation（例如把 \(N_\t{DP}\) 和 \(N_\t{EP}\) 合称 DP，此时其实际意义是同时在运行的 micro-batch 总数）。

每 \(N_\t{EP}\) 张卡分一组，做普通的 ZeRO DP：

  • attn 层正常 DP；
  • MoE 层：
    • 参数方面，每张卡只分到 \(E/N_\t{EP}\) 个 expert，但拥有（或通过 ZeRO-3 假装拥有）完整的 router。
    • 计算方面，每个 token 先过 router 得到去向，然后 a2a 路由到持有相应 expert 的卡，expert FFN 后再 a2a 路由回来。

MoE 的一个核心痛点是保证负载均衡，但局部 router 并不知道这一点，因此可能出现不同节点收到不同数目的 token。有一些神秘的 routing 方法，诸如：

  • 让 Expert 自己去选 token。
  • 在 Expert 和 token 间跑 Sinkhorn。

但这些方法没一个在 EP 中能要的。因此一般采用以下几者的结合：

  • 首先在训练时用一个 loss 强制负载均衡，且在训练初期甚至要额外强调这个 loss。
  • 然后在训练时一般为所有 expert 设置容量上限，超过上限的 token 会被丢弃，也就是 dropout 掉，MoE 层变成 identity（因为有残差连接）。（显然，这意味着手动 dropout 是有必要的，不然下一层压根没见过 identity 就会出现 ood）。推理时为了速度可以保留容量上限，也可以为了质量不设上限（不过理论上你训完了不应该已经负载较为均衡了吗？）
  • 还可以通过给 router 分数中加可学习的噪声防止塌陷、使用公共专家、分层路由（先 route 到节点再在内部 route 到具体专家）、动态监测专家侧的流量并回输给 token 侧调整权重等方法。

对于 MoE 层，有时我们会更倾向于 EP 而不是 TP，出于以下几种可能的原因：

  • EP 没有把矩阵切细，对 GeMM 更友好。
  • EP 需要的通信量更小（特指 \(k\leq2\) 时；如果 \(k\) 更大结论会反过来）
  • EP 可以分块进行，很容易重叠不同 token 的 routing、通讯和 Expert 运算；但 TP 的通讯是严格全阻塞的。
  • 当 \(E=N_\t{EP}\) 时，EP 不需要在 expert 端再次 routing，效率能再上一个台阶。

II. Detailed Comparisons

本节从各种角度对比上述并行方案。作出以下简写：

  • OS：优化器状态。
  • G：梯度。
  • W：模型权重。
  • A：激活值。

II.I. Partitioning Axes

DP：从 batch 角度切分。ZeRO-1/2/3 分别消除了 OS/G/W 的副本式冗余，但无法切分 A，以至于 ZeRO-3 中 A 会成为瓶颈。只用 DP 的场景，ZeRO-R 并不能真的减少单卡 A 的显存。

PP：从 layer 角度切分，同时作用于 OS/G/W/A。必须使用较大的 GA step 以减少 bubble，同时具体的流水线调度方案也会与峰值 A 开销有关。

TP：对于 FFN 层，从矩阵角度切分；对于 attn 层，从 head 角度切分。同时作用于 OS/G/W/A。

SP(Ulysses)：对于 FFN 层，从 sequence 角度切分，此时和 DP 几乎一致，因此依赖于具体 ZeRO 方法；对于 attn 层，同 TP，此时同时作用于 OS/G/W/A。

CP(Ring Attn)：FFN 层是 DP；attn 层因为深入到具体 flash-attn 内部，相当于隐式的 AC，所以不好说具体咋切的。

EP：默认与 DP 共用，所以首先有 batch 角度切分，同时伴有对 expert 切分。非 MoE 层的分析依赖于具体 DP 模式，MoE 层则是对 OS/G/W/A 的协同切分。

II.II. Communication Approaches

本节默认训练。

DP：

ZeRO-0

graph LR A[FP] --> B[BP] --> C[梯度 A-R<br/>2δΦ Bytes]--> D[全量优化器更新] --> E[全量参数更新]

其中 BP 与 A-R 可分块重叠（算完一块就立刻 All-Reduce），是最常见的 ZeRO-0 并行方法；A-R 与后两步也可分块重叠（Reduce 完一块就立刻更新优化器与参数），但因为瓶颈往往在计算，所以一般朴素的 ZeRO-0 不会干这个，但是 Megatron-LM 等框架是支持后两步的自动重叠的。

通信频率是 per-step；启用分块重叠就是 per-bucket。

ZeRO-1

graph LR A[FP] --> B[BP] --> C[梯度 R-S<br/>δΦ Bytes]--> D[分片优化器更新] --> E[分片参数更新] --> F[参数 A-G<br/>δΦ Bytes]

其中 BP 和梯度 R-S 仍然可以分块重叠，且是最主要的重叠；如果该节点的分片梯度准备好了，也可以立刻开始优化器/参数更新乃至最后的 A-G，但仍然不主流。

通信频率仍是 per-step 或 per-bucket。

ZeRO-2

graph LR A[FP] --> B[BP] <--"交替"--> C[梯度 R-S<br/>δΦ Bytes]--> D[分片优化器更新] --> E[分片参数更新] --> F[参数 A-G<br/>δΦ Bytes]

BP 和梯度 R-S 的分块重叠不是可选项而成为必须项（毕竟你不能等 BP 全跑完再 R-S，那就白进行梯度分片了）。分片梯度也可以立刻启动之后流水线。

通信频率是 per-bucket。

ZeRO-3

graph LR A[FP] --> B[BP] <--"交替"--> C[梯度 R-S<br/>δΦ Bytes]--> D[分片优化器更新] --> E[分片参数更新] A <--"交替"--> G[参数 A-G<br/>δΦ Bytes] B <--"交替"--> H[参数 A-G<br/>δΦ Bytes]

通讯-计算重叠主要包括：

  • FP 和参数 A-G 的重叠，也即参数 prefetch。只要 prefetch 足够多且延迟不过大，则可以被完美掩盖。
  • BP 和参数 A-G 的重叠，同上。
  • BP 和梯度 R-S 的重叠，同 ZeRO-2。只要缓冲 bucket 足够大且延迟不过大，则可以被近似完美掩盖。
  • 分片梯度 Reduced 完毕后立刻开始优化器/参数更新，可选项。
  • 通信频率是 per-bucket + prefetch。

PP：

  • 通讯全是 P2P 的，而且都是非阻塞，很容易掩盖几乎所有通讯代价（除了预热时的前几个 micro-batch 有 不可避免的 bubble），而且不同 micro-batch 在 FP 中、BP 中、FP-BP 之间都可以互相掩盖。
  • 如 Zero-Bubble PP，单个 pipeline 内部在 BP 时可进一步细分为 A 和 G 两部分，A 准备好就可以直接发射，G 则可以慢慢算。
  • Interleaved 1F1B 因为拆的更碎，虽然通讯更频繁，但掩盖通讯会更容易。
  • PP 在 ZeRO 流水线中对应 梯度计算 部分；如果和 ZeRO 联合使用，则如之前分析，可以在计算完后立刻启动 R-S，R-S 完后立刻启动优化器更新和参数更新。
  • 如果使用 AC，则 AC 的 Forward 计算也可以用来掩盖通讯。
  • 如果需要 CPU offload，也可以与计算/通讯重叠。
  • 但作为流水线，最大的问题还是 牵一发而动全身，单个任务的延迟很容易卡住整个流水线，也因此更加在意 传输延时。
  • 总通信量：在每个 pipeline 边界上，FP + BP 各需发送 \(\delta bSH\) 的激活值。
  • 通信频率是 per-border。

TP：不论是 FFN 还是 attn 都遵循以下模式

graph LR A[输入] --> B["(identity,all-reduce)"] --> C["内部矩阵乘法/attn+proj"] --> D["(all-reduce,identity)"] --> E[输出]
  • 阶段二和阶段三（仅在 BP 时）、阶段三和阶段四（仅在 FP 时）都可以分块重叠的，也即传输完一个小 tile 就立刻开始走内部流程，算完再立刻传出去。
  • 但一方面会把 GeMM 切得稀碎，降低原生内核效果；另一方面还需要流水式的 GeMM 和 all-reduce，甚至必须要求不同节点的 tile 匹配。
  • 总体而言实现很麻烦，因此 默认的 Megatron-LM 不搞这个，不过 近年来很多新型框架会搞。
  • 总通讯量：每一个 ffn 或 attn block 都在 FP 的输出前要一次 all-reduce、BP 的输入前要一次 all-reduce，而每次 all-reduce 的总流量是 \(2\delta bSH\)。因此是 \(4\delta bSH\) per-block、\(8\delta bSH\) per-transformer layer。
  • 通信频率是 per-block。

SP(Ulysses)：

graph LR A[输入] --"QKVproj(ZeRO)"--> B[QKV] --"a2a"--> C["head 内部 QKV"] --"self-attn"--> D["softmax(QK^T)V"] --"a2a"--> E["tokenwise-latent"] --"proj(ZeRO)"--> F["attn 输出"] --"FFN(ZeRO)"--> G["FFN 输出"]

能重叠的部分包括：

  • 几个让 ZeRO 管理的模块调用 ZeRO 内部的重叠方式。
  • QKV proj 与 a2a 通讯的重叠：算好一块就开始 a2a。
  • a2a 与 attn 的重叠：传好一个 head 就开始 attn。
  • attn 与第二次 a2a 的重叠：算完一个 head 的 attn 就开始 a2a。
  • a2a 与 proj 的重叠。
  • 标准实现一般不包括后面几个 a2a 重叠，直接使用阻塞式 a2a；新兴框架可能会用。
  • 总通信量：两次 a2a 一次要传输共计 \((2N_{kv}+N_q)bSd/N_\t{SP}\) 的 QKV，一次要传输 \(bSH/N_\t{SP}\) 的 attn 输出；因此 a2a 传输量是 \(2\delta bS((2N_{kv}+N_q)d+H)/N_\t{SP}=4(\dfrac{N_{kv}}{N_q}+1)\delta bSH/N_\t{SP}\)。
  • 通信频率是 per-block + ZeRO。

CP(ring attn)：

graph LR A[输入] --"QKVproj(ZeRO)"--> B[QKV] --"ring-attn"--> E["tokenwise-latent"] --"proj(ZeRO)"--> F["attn 输出"] --"FFN(ZeRO)"--> G["FFN 输出"]

主要的重叠在于 ring 通信的重叠，不过 QKVproj 和 ring、ring 和 proj 也能重叠一部分。

总通信量：FP 时要把 KV 转一圈因此是 \(2\delta bSN_{kv}d\)；BP 时，因为 flash-attn 要重计算，所以 KV 和 dK dV 要各绕一圈。因此除 ZeRO 外的总通信量为 \(6\delta bSN_{kv}d=6\dfrac{N_{kv}}{N_q}\delta bSH\)。

通信频率是 per-block + ZeRO。

可以发现，Ulysses 比起 ring-attn，通信量是 \(\Theta(1/N_\t{SP})\)，理论上会随着并行度上升而下降；但代价是其拓扑必须支持 a2a 和高并发，而且 a2a 的重叠显然比 ring 要困难很多。ring-attn 相当于 用通信量换拓扑结构。

EP：

graph LR A[输入]--attn(ZeRO)-->B[attn 输出] --> C["router(ZeRO)"] --"a2a"--> D[目标 expert] --> E[expert FFN] --a2a--> F[输出]

主要的重叠在于 attn 至 router 这条 pipeline 与 a2a 通讯的重叠，以及 expert FFN 与 a2a 通讯的重叠。

一来一回两次 a2a，总通信量为 \(4\delta K_\t{MoE}bSH\)。通信频率是 per-block + ZeRO。

方法 通讯类型 重叠方式 关键要求 总通讯量 (FP+BP) 通信频率
ZeRO-0 A-R 分块重叠 带宽优先
适合 inter-node		 \(2\delta\Phi\) per-step / bucket
ZeRO-1 R-S + A-G 分块重叠 带宽优先
适合 inter-node		 \(2\delta\Phi\) per-step / bucket
ZeRO-2 R-S + A-G 分块重叠 带宽优先
适合 inter-node		 \(2\delta\Phi\) per-bucket
ZeRO-3 R-S + A-G 分块重叠 + 预取 带宽优先
延迟也有需求
不太适合 inter-node		 \(3\delta\Phi\) per-bucket + prefetch
PP p2p pipeline 延迟优先
非常适合 inter-node / cluster		 \(2\delta bSH\)
(per-border)		 per-border
TP A-R 分块重叠 (罕见) 带宽延迟均高要求
几乎必须 intra-node		 \(4\delta bSH\)
(per-block)		 per-block
SP(Ulysses) ZeRO + a2a ZeRO + a2a 带宽优先
一般 intra-node，inter 需要特殊高带宽		 \(4(\dfrac{N_{kv}}{N_q}+1)\delta bSH/N_\t{SP}\)
(a2a-only)		 per-block + ZeRO
CP(ring attn) ZeRO + p2p ZeRO + ring 带宽优先
非常适合 inter-node		 \(6\dfrac{N_{kv}}{N_q}\delta bSH\)
(ring only)		 per-block + ZeRO
EP ZeRO + a2a ZeRO + a2a 带宽优先
一般 intra-node，inter 需要特殊高带宽		 \(4\delta K_\t{MoE}bSH\)
(a2a only)		 per-block + ZeRO

可以发现，一般用标准通讯操作（R-S + A-G）的，都可以接受 inter，除了 TP（因为没法重叠）；所有用 p2p 的都非常适合 inter；所有用 a2a 的如果硬要上 inter 就需要特殊设计。

II.III. Computational Details

本节讨论若干计算细节。

  • 是否可能导致矩阵过小，GeMM kernel 跑不满？

    • DP/PP：可能，但一般是间接的。\(B=b\times m\times N_\t{DP}\times N_\t{EP}\)，如果 \(b\) 太小 GeMM 就可能退化为 GeMV；当 \(B\) 不变而 \(N_\t{DP}\) 过大或 \(m\) 过大（EP 用 GA 来掩盖 bubble）就会导致 \(b\) 过小。
    • TP：非常可能，尤其是若模型本身就不大，\(H\) 比较小，则切完后就更小了。所以 TP 一般不会搞太多。
    • SP(Ulysses)：罕见。如果 SP 把序列切太细可能导致 GeMM 退化，但一般启用 SP 都是在序列真的太长、单卡塞不下的场景，也不会切得太过分。
    • CP(ring-attn)：可能。除了 FFN 阶段退化以外，如果切太细则不好流水线掩盖通讯。
    • EP：非常可能，尤其是负载不均衡的场合。

  • 是否可能退化为 communication-bounded？（其实和上一节的通信掩盖分析有部分重叠）

    • DP：ZeRO 0/1 因为通信是 per-step 的所以几乎不可能，2 只要开大 bucket 也不太可能，3 有可能，通过开大 prefetch 可以缓解。
    • PP：pipeline 所以几乎不可能。
    • TP：非常可能，因此几乎不出节点。
    • SP(Ulysses)：有可能，尤其是跨节点 a2a 通讯环节。
    • CP：不太可能，ring 还是太容易掩盖了。
    • EP：非常可能，routing + a2a 高度不可预测、容易拥堵。

  • 是否仍可直接使用 fused kernel？

    报个菜名。这些 kernel 可以以下分类：

    • attn 内部算子： flash-attn、fused attn-scale（除以 dim scale）+ mask + softmax 等。
    • token-wise 算子（每个 token 独立）：Fused MLP（字面意思）、FP8 fused kernel（量化 + 计算 + 反量化）、fused layernorm / RMSnorm、fused bias + dropout + residual add 等。
    • MoE 算子：fused router（字面意思，一步到位算 MoE router）、fused dispatch（根据 router 结果进行 permute + 发送）、fused combine（dispatch 的反向操作）等。
    • 其它算子：fused RoPE、fused optimizer、fused all-reduce + residual add 等。
    • DP / PP：完美适配所有上述 kernel，因为这些 kernel 基本上都局限于 layer / block 内部，没有跨 layer 的；有一些罕见的跨 layer 的 fusion 可能在 PP 中不适用。
    • TP：会破坏涉及 latent 的算子，比如 MLP 和 layernorm，这些东西必须用 TP 特制版；但不会破坏 attn 内部算子。
    • SP (Ulysses)：完美适配 attn 内部 / token-wise / MoE 算子相关，但 RoPE 这种需要 global offset 的需要微调。
    • CP (ring-attn)：需要特制 attn 算子；token-wise 兼容；RoPE 要微调。
    • EP：需要特制 MoE 算子，其它兼容。

  • 是否容易负载不均衡：

    • DP：ZeRO-0 所有卡地位平等；1 的 OS 一般可以完全均分；2 的 Reduce-Scatter 也相对平滑；3 因为参数分配不均衡，可能会出现通信不对称、某些卡的 OS 算得慢、显存碎片化、不同层 Reduce-Scatter 时机不同等一堆隐蔽的问题。
    • PP：比较容易，一方面因为架构不对称 (Embedding 与 LM Head) 会有不同负载，另一方面启动和收尾阶段必然有部分卡闲置，因此需要手动调整。
    • TP：极不可能，因为矩阵 / head 的切分是粗暴且均匀的。
    • SP (Ulysses)：极不可能，因为切分是均匀的。
    • CP (ring-attn)：在存在 causal-mask 的场景，负载会显著不均，此时需要使用非常复杂的调度方法。non-causal 的场合则不太可能。
    • EP：非常可能，所以要特殊处理。

II.IV. Bounds & Scales

本节分析这些算法支持的并行度，包括并行度上限和达到上限前能否稳定 scale。

DP：

  • 上限：\(B=b\times m\times N_\t{DP}\times N_\t{EP}\)，而 \(B\) 作为 optimizer 看到的全局 bsz，受模型、优化器等限制。\(B\) 有上限则 \(N_\t{DP}\) 有上限，同时在到达上限之前就可能因为 \(b\) 过小而 scale 减缓，正如上一节所述。
  • scale：ZeRO 0/1 能稳定 scale 且容易掩盖，但 ZeRO 3 因为要 gather 权重，所以在极大规模时 scale 有限。

PP：

  • 上限：受模型层数限制，而且必须保证切分均匀，不然会被最慢的一个卡死。
  • scale：困难，因为 scale 越大 bubble 越大，bubble 越大就要拉大 GA step \(m\)，这又受 \(B\) 限制。
  • 实践：一般 4-8 就够了。

TP：

  • 上限：受节点卡数和 \(N_{kv}\) 双重卡死。
  • scale：没法跨节点。
  • 实践：单节点一般 8 卡，所以只能取到 8。

SP(Ulysses)：

  • 上限：受限于 \(N_{kv}\)。
  • scale：a2a 受限于具体拓扑，跨交换机 scale 困难。
  • 实践：\(\min(N_{kv},\t{单交换机卡数})\)，且一般对 32K-128K 等中等长度文本使用，文本长度更大就要上 ring-attn 了。

CP(ring-attn)：

  • 上限：只要 context length 足够长就可以任意切。
  • scale：ring 还是太权威了，稳定且容易重叠。

EP：

  • 上限：expert 数。
  • scale：网络上总的 a2a 流量越来越大、负载越来越不均匀，因此总体而言是糟糕的。

II.V. Agnostic to Implementation

并行是否能真的做到 agnostic，以至于上层来看就像一张大卡一样？

DP/PP：严格等价，但是如果在低精度场景，规约顺序会产生数值截断噪音，在极大规模场合可能会累计误差，但一般不影响收敛。

TP：只要框架写得对就等价，但 dropout 时如果所有卡持有同一个种子，会得到相同的 dropout mask，这显然是不合适的；因此 Megatron 等框架会魔改 RNG，使得不同卡的 mask 拼接起来刚好等于理想单卡的 mask，保证 agnostic。

SP/CP：严格等价，甚至有赖于 flash-attn 的数学性质，连规约顺序都可以保证相同！

EP：不等价，因为有丢弃、随机噪声、路由 loss 等一堆东西，强迫 MoE 算法设计者必须关注具体并行细节。

III. Parallelisms With Others

本节 不会 讨论所有并行方法的两两组合，那样低效且不符合真实工程经验。取而代之，会建立一个全局且通用的框架；常见的并行方案则往往只会取该框架的若干子维度。

III.I. Process Grid

前文中我们一直使用 \(N=N_\t{TP}\times N_\t{PP}\times N_\t{CP}\times N_\t{DP}(\times N_\t{EP})\) 的分解。这里提供一个更通用的框架：

\[N=\sum_{i=1}^{N_\t{pipeline}}N_\t{intra}^{(i)}\times N_\t{inter} \]

其中：

  • \(N_\t{pipeline}\) 就是前文的 \(N_\t{PP}\)，所有卡在 pipeline 上被组织成多少组。它不一定真的等于总 pipeline stage 数，尤其是 interleaved 1F1B 这种场合会有 stage 数大于 \(N_\t{PP}\)。
  • \(N_\t{intra}^{(i)}\) 指第 \(i\) 组 PP 卡，或者说第 \(i\) 道工序（假如不用 interleaved 的话）中，服务每个 micro-batch 的卡数。注意这里特别强调了每一组的卡数可能是 不同 的，也就是所谓的 hetero pipeline：尽管主流框架一般要求所有工序都使用相同数目的卡，如果发生负载不平衡应当切 pipeline 而不是分更多卡，但某些前沿探索可能会使用这种模式。同时，每一道工序，甚至工序中的不同层在内部都可能有不同的组织方式，例如在 attn 层用 TP 进行 head-wise 分割，但在 MoE 层混用 TP + EP 等。前文框架中的 TP、SP(Ulysses)、CP(ring-attn) 和 EP 的 MoE 层都与这部分有关。
  • \(N_\t{inter}\) 指有多少个 micro-batch 被同步运行。前文的 DP 和 EP 都会贡献 micro-batch。

如果不考虑 hetero pipeline，则可以简写为 \(N=N_\t{pipeline}\times N_\t{intra}\times N_\t{inter}\)。

III.II. Disentanglement Between DP and ZeRO and EP

首先将 DP 和 ZeRO 解耦。

  • ZeRO 是一种 通用的、维护同一份参数在多张卡上副本 的方式；不管哪种 ZeRO 都可以维护副本的跨卡逻辑一致性。当然也存在一些 非 ZeRO 的副本组织方式。本文用 参数副本 来代指这种现象。
  • DP 因为要保证所有 batch 都在同一个模型上跑，所以是 产生参数副本的最常见场景；同时，如果要搞一些 token-wise invariant 的操作比如说 FFN CP/SP，因为所有 token 都要通过同一个 FFN，同样会产生参数副本。
  • 因此，参数副本组是一个大概念，可能横跨整个 DP/EP 以及 FFN 层的 SP/CP。

而脱离 DP 语境后，ZeRO 其实也很简单：

  • ZeRO-0：在 step 前确保所有副本的梯度被 all-reduce 了即可。
  • ZeRO-1：不 all-reduce，而是在 step 前把梯度 reduce-scatter 到分片卡，分片卡本地更新分片优化器和参数，再 all-gather 回来。
  • ZeRO-2：确保所有梯度在产生后立刻被 reduce-scatter 到对应卡。
  • ZeRO-3：把副本分片存储，需要时再 all-gather 过来。

不同的副本组可以使用不同的同步模式；同一个副本组内部还可以混合多种模式，例如对于跨机柜分布的 ZeRO 组，在机柜内部用 ZeRO-3 但机柜间退化为 ZeRO-1。PyTorch FSDP 等可以显式把这个分割以 mesh 形式组织。

至于 ZeRO-R，如果把它放在 DP 中就是削足适履了，它除了引入一些微不足道的激活值负载均衡没有任何意义，还白增了很多通讯。真正重要的场景比如说 TP，其输入激活值就完全是副本，此时可以用 ZeRO-R，这才是其大显身手的地方。但是注意，正如上述分析，ZeRO-R 不适合跨 DP 组同步，因为它们的激活值压根不同。它与 AC 是协同关系：AC 决定哪些激活值或辅助量要被 save for BP，而 ZeRO-R 决定怎么压缩。

进一步，如果把 EP 放到 DP 的框架下，会发现：

  • 理论上来讲，MoE 层可以正常 DP，所有 batch 都持有 MoE 层的完整副本，并可以通过 ZeRO 进行切分和分片存储。
  • EP 同样有参数副本分散存储的行为，但不像 ZeRO3 一样分片，而是 以 expert 为单位分散。
  • 另一个区别是，ZeRO-3 是 Weight-to-Data，从存有分片权重的卡上拉来权重，然后在存有激活值的卡上计算得到下一步激活值；
  • EP 是 Data-to-Weight，把激活值发射到存有权重（完整的 expert）的卡上计算，算完后再把激活值拉回来。
  • 整个参数副本组可以混用 ZeRO 和 EP，一个 MoE 层可以以 expert 为单位拆分到各个卡上，同时还可以复制多个副本并用 ZeRO 管理。而且理论上 EP 可以在 任意 MoE 层的参数同步组 进行，不过实际实践时有一些细节需要打磨。

以下有一个对比表格，但需要注意 ZeRO-R 和另外两者的性质有根本区别：ZeRO-R 的生命周期是 FP 至 BP 间，是为了处理 AC 而使用的；而另外两者的生命周期贯穿整个训练。不过它确实是持有权重的 TP 卡在从组里其它卡拉取数据，一定程度上可以称作 Weight 方发起的 D2W。

方法 性质 发起方
ZeRO-3 W2D D
ZeRO-R D2W W
EP D2W D

III.III. Block Layout Contract

如果不同层要使用不同的 intra-micro-batch 划分方式，那势必需要约定一些传输接口；如果不满足接口，必须像 SP(Ulysses) 一样进行 a2a 的 layout transform，引入额外的开销。

这个 layout 几乎总是遵循下述形式：每处 border 的激活值都服从

\[\Big((b,S/N_\t{part},H)\times N_\t{part}\Big)^{N_\t{rep}} \]

的形式，其中：

  • \(b\) 是 micro-bsz。
  • \(N_\t{part}\) 衡量 context 角度被切成了几段，也就是之前框架中的 \(N_\t{CP}\) 或 \(N_\t{SP}\)。特别地，在某些框架（如 Megatron-SP），TP 也会贡献 \(N_\t{part}\)，将在下文详细阐明。
  • 定义 \(s=S/N_\t{part}\)。
  • \(H\) 是完整的 latent dim。
  • \(N_\t{rep}\) 意味着这个 border 处的激活值有 \(N_\t{rep}\) 个副本，一般对应之前框架中的 \(N_\t{TP}\)，但也不一定。
  • 始终满足 \(N_\t{intra}=N_\t{part}\times N_\t{rep}\)，但不同层的拆分可能不同。

处于同一个 TP 组的所有卡，在同一模块：

  • FP 时接受相同的输入端激活值，给出相同的输出端激活值（通过 all-reduce 实现）；
  • BP 时接受相同的输出端梯度，给出相同的输入端梯度（通过 all-reduce 实现）；
  • 这些卡持有的权重是被分片的。
  • 通常不希望所有卡对同一个激活值进行相同的操作，这样的计算是重复的；但有时为了减少通信量，也会不得已这么做。因此比如说 attn 最后的 proj，应当是每个卡各自把 \((b,S,H/N_\t{rep})\) 的局部激活值升维到 \((b,S,H)\) 然后再 all-reduce。
  • TP 相当于给整个模块做了一个 包装，包装内部的宽度减少了：attn 头数缩水 \(1/N_\t{rep}\)、FFN 宽度缩水 \(1/N_\t{rep}\)，但内部看来会是一个完整且自洽的瘦模型，只需要在边界处做好维护即可。
  • all-reduce 可以被细分并进行算子融合，将在下文分析。

在 border 处可以切换 TP + CP/SP 组合。这里的 border 指任意模块的交界处，可以是 PP 交界处也可以不是；但出现切换就需要 layout transform 对齐接口。

现在假设已经被 TP 包装好了，也即在当前的卡看来，内部 attn 头数其实是 \((N_q/N_\t{rep},N_{kv}/N_\t{rep})\)（如果出现不完全切分可能多卡要共用同一个 KV-proj，此时构成参数副本）；FFN 宽度其实是 \(4H/N_\t{rep}\)，而且默认已经在模块头尾做好了 all-reduce。则内层的 SP/CP 可以 agnostic to 外界的 TP，只需要分析内层行为。

对于 attn 层：

  • SP(Ulysses)：
    • 持有 \((b,s,H)\) 的输入激活值。
    • （BP 时需要 TP all-reduce）
    • 用 QKV proj 得到 \((b,s,N_q/N_\t{rep},d)\) 的 Q 和 \((b,s,N_{kv}/N_\t{rep},d)\) 的 KV。在同一个 SP 组中，QKV proj 构成参数副本。
    • layout transform 为 \((b,S,N_q/N_\t{intra},d),(b,S,N_{kv}/N_\t{intra},d)\) 的 head-wise 独立格式。
    • 同上，不完全切分时，共用的 KV head 构成参数副本。
    • 内部 head-wise attn 变成 \((b,S,N_q/N_\t{intra},d)\) 的输出。
    • reshape 到 \((b,S,H/N_\t{intra})\)。
    • 逆 layout transform，回到 \((b,s,H/N_\t{rep})\)。
    • proj，变成 \((b,s,H)\)，构成参数副本。
    • （FP 时需要 TP all-reduce）
  • CP(ring-attn)：
    • 持有 \((b,s,H)\) 的输入激活值。
    • （BP 时需要 TP all-reduce）
    • 得到 \((b,s,N_q/N_\t{rep},d)\) 的 Q 和 \((b,s,N_{kv}/N_\t{rep},d)\) 的 KV，QKV-proj 构成参数副本。
    • ring-attn，得到 \((b,s,N_q/N_\t{rep},d)\) 的输出，然后 reshape 到 \((b,s,H/N_\t{rep})\)。
    • proj 到 \((b,s,H)\)，构成参数副本。
    • （FP 时需要 TP all-reduce）
  • 此外还有一些不常见 attn 框架，不过大体是与 contract 兼容的。

对于 FFN 层：

  • SP/CP：
    • 持有 \((b,s,H)\) 的输入激活值。
    • （BP 时需要 TP all-reduce）
    • 过标准的 \((b,s,H)\to(b,s,4H/N_\t{rep})\to(b,s,H)\) 模式，且同一个 SP/CP 组中的 proj 构成参数副本。
    • （FP 时需要 TP all-reduce）

对于 MoE 层：

  • SP/CP + EP：
    • 持有 \((b,s,H)\) 的输入激活值。
    • （BP 时需要 TP all-reduce）
    • 过 router，得到 \((b,s,K_\t{MoE})\) 的 routing 目标。router 本身在整个 SP/CP 组中构成参数副本。
    • 然后是 MoE。整个 SP/CP × DP/EP 构成一个大的同步组，里面所有 token 都经过 同一个 MoE layer，因此 EP 可以不局限于 inter 维，在整个同步组里自由分拆 expert 或创建 expert 副本并 ZeRO 同步。有一些专业术语：
    • expert placement group 指整个 MoE layer 的所有 expert 被拆到了哪些卡上，组内所有卡的 expert 集合不交；
    • same-expert replica group，同一个 expert 被 ZeRO 式地复制了多少份副本。
    • token dispatch group，这是 token 端的概念，哪些卡上的 token 要共用同一个 expert placement group。
    • 在 expert 端，跑标准的 \((b,s,H)\to(b,s,4H/N_\t{rep})\to(b,s,H)\) 操作。
    • （FP 时需要 TP all-reduce）

现在来玩一点花样。比如说，TP 结尾都有一个 all-reduce，而这个 all-reduce 可以被拆成 reduce-scatter + all-gather；两者之间可以额外插入一些 token-wise 的操作，比如说 layernorm 或 dropout；这种东西原本会被所有 TP 卡重复计算。

因此，有 Megatron-SP 这种东西，它实际是一种 TP：

  • 在模块最后，把 \((b,s,H)\) reduce-scatter 到 \((b,s/N_\t{rep},H)\)。
  • 进行 layernorm / dropout。pre-norm 或 post-norm 均可，因为在此阶段跨过了 border。甚至该 border 可以跨越 PP，这样 PP 的通信量还能更小，不用重复发送所有 TP 副本。
  • 在下一个模块的开始处，把它 all-gather 到 \((b,s,H)\)。
  • 它的思想和 ZeRO-R 有点类似，同样是把冗余的激活值分散处理，不过 ZeRO-R 是一种存储策略，而 Megatron-SP 是一种类似算子融合的技巧。

另一种理解方式是，Megatron-SP 在 border 处的 contract 切换为 \(N_\t{rep}=1\)，\(N_\t{part}=N_\t{intra}\)，但是相关的 layout transform 和 all-reduce 融合了。

III.IV. Pipelines

在上面这套框架下，PP 的任务是清晰的：

  • 每个 border 上，本应有一个 layout transform；只不过当前后 layout 不变 且 在同一组卡上进行，可以省去这个 transform。
  • PP 让 border 两侧 在不同卡上进行，此时即需要显式、强制的传输。如果 border 两侧的 layout 相同，通讯是严格 p2p，否则还要伴随 a2a。

除此之外，PP 的若干问题，比如说 schedule、bubble、weight tying 等问题，都是 PP 内部的问题，与其它东西相对解耦。

III.V. Conclusion

所以可以做以下总结：

  • DP/EP：把参数在 batch 维创建很多副本。任何 ZeRO 同步组都可以用 ZeRO + EP 处理；划分比例任意，不限制在 batch 维。内层的框架可以适当对这些副本 agnostic，但 MoE 路由负载均衡相关问题还是需要关注的。
  • TP：把激活值创建很多副本，在相邻两次同步间这些副本彼此独立且 agnostic，只需要在同步时 all-reduce 即可。也可以把 all-reduce 拆开来进行适度的算子融合。
  • FFN/MoE 时的 CP/SP：在 context 维切分，切分出的每一块彼此独立，输入相同的、靠 ZeRO/EP 同步的 FFN/MoE 模块。
  • attn 时的 SP(Ulysses)：a2a layout transform 后 head-wise attn 再 a2a 回来 proj。
  • attn 时的 CP(ring-attn)：直接 ring-attn 然后 proj。
  • PP：显式把 border 拆到不同卡上，强制 p2p 传输，并伴随可能的 layout transform。
posted @ 2026-07-03 13:08  Troverld  阅读(15)  评论(1)    收藏  举报