ラブパラ

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

O. Introduction

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

O.I. Notation

本章将所有 notation 统一。

并行相关配置：

$N$：总卡数。
$N_\t{TP}$：张量并行 (切分矩阵乘法) 数。
$N_\t{PP}$：流水线并行 (切分模型层结构) 数。
$N_\t{CP/SP}$：上下文并行 (切分输入序列长度) 数。需要注意，这里同时提到了 CP 和 SP，这两者有一些隐式的区别，本文虽然会遵循惯例区分二者，但读者如果不想区分也没有任何关系。这两者并不是非此即彼的关系，也可以同时启用；下文中的 notation 可能有适当混用，请自行区分。
$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$。

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

这样做并没有省计算量，反而引入了额外的传输开销，因此只有一个 micro-batch 时，PP 不能提升带宽，但可以减少每张卡的参数存储开销。只有当多个 micro-batch 同时在流水线时，才会出现真正的、典型的流水线重叠行为；而在训练态这只能依靠 Gradient Accumulation 实现。

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 就能胜任。代价是需要更多传输。

但是，此时一张卡上持有的多个 virtual stage 需要约定一些执行顺序：如果是「随到随算」的模式，虽然调度简单，但切换开销太大，而且 L2 cache 不能被很好利用，所以一种常见方式是 Grouped Processing，一张卡连续处理 $g$ 个 micro-batch 后再切换到下一个 virtual stage。

为了保证 1F1B 可以达到稳态，常见做法要保证 $g\geq N_\t{PP}$，这样一个任务在流过其它卡各一轮后，刚好能回到初始卡并切换 stage；为了保证收尾时不出现空转，常见做法要保证 $m\bmod g=0$ 或 $m\bmod g\geq N_\t{PP}$，理论同上。$g$ 越大，则 context switch 开销越小，但 warmup 时长越长，同时 bubble 可能变大，且单个任务在 pipeline 中停留时间也变长（这一点在推理时格外重要）。

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

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

前述所有 PP 都是 同步 PP，强制所有 micro-batch 都看到同一组参数，因此需要 GA，也即定期清空 pipeline 进行 step，两次 flush 间则完全没有参数更新。也存在一些 异步 PP，允许不同 micro-batch 看到不同参数；但这样会导致上层算法不再对 PP agnostic，不能只看一个全局 bsz $B$ 就下结论。因为过于复杂，并非本文重点。

I.III. Tensor Parallelism

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

其中 (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 只会选择在行或列切一次，而不会都切；此时最多支持连续两次矩阵乘法（第一次竖着切，第二次横着切）。但是，也存在一些特种 TP，比如连续进行超过两次的，或者如同 GPU 内部的 GeMM 一样行列都切；这种 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["Q1:(b,N_q/N_TP,S,d) KV1:(b,N_kv/N_TP,S,d)"] X2 --GeMM--> QKV2["Q2:(b,N_q/N_TP,S,d) KV2:(b,N_kv/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]

此外，embedding 和 LM head 的 TP 也需要特殊实现，被称作 Vocabulary Parallelism。

embedding 层：
- 在单卡场景是对 $(b,S,V)$ 的 one-hot 输入序列右乘 $(V,H)$ 的 embedding 矩阵得到 $(b,S,H)$ 的激活值。
- 现在在多卡场景，则每张卡持有 $(V/N_\t{TP},H)$ 的 vocabulary 分块 embedding 矩阵；与输入的 $(b,S,V)$ 相乘时使用 masking 机制，只有 one-hot 项落入这张卡的局部 vocabulary 时才会贡献结果。然后过一个 all-reduce 把所有卡计算得到的局部 $(b,S,H)$ 相加得到全局 $(b,S,H)$。
LM 层：
- 在单卡场景是对 $(b,S,H)$ 的激活值右乘 $(V,H)$ 的 embedding 矩阵（一般与 embedding 层共享）的转置，得到 $(b,S,V)$ 的 logits，然后算 CE loss。
- 在多卡场景，仍然让每张卡持有 $(V/N_\t{TP},H)$ 的分块矩阵，则乘以转置后得到 $(b,S,V/N_\t{TP})$ 的局部 logits。此时通过轻量级的通讯，可以得到：
  - 局部最大 logit。
  - 全局最大 logit（第一次 all-reduce）。
  - 局部 expsum。
  - 全局 expsum（第二次 all-reduce）。
  - 所有卡都得到相同的全局 log-sum-exp。
- target 标签以 mask 的形式作用于 log-sum-exp 后的 logprob，每张卡得到局部的 loss，最后通过一个 all-reduce 聚合出全局 loss；或者直接在反向传播时累加梯度。

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，其 drop-rate 需要与 capacity 导致的 drop-rate 协调。
也存在一些 dropless 的方法。
还可以通过给 router 分数中加可学习的噪声防止塌陷、使用公共专家、分层路由（先 route 到节点再在内部 route 到具体专家）、动态监测专家侧的流量并回输给 token 侧调整权重等方法。

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

EP 没有把矩阵切细，对 GeMM 更友好。
EP 需要的通信量更小（特指 $K_\t{MoE}\leq2$ 时；如果 $K_\t{MoE}$ 更大结论会反过来）
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 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 δΦ Bytes]--> D[分片优化器更新] --> E[分片参数更新] --> F[参数 A-G δΦ Bytes]

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

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

ZeRO-2

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

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

通信频率是 per-bucket。

ZeRO-3

graph LR A[FP] --> B[BP] <--"交替"--> C[梯度 R-S δΦ Bytes]--> D[分片优化器更新] --> E[分片参数更新] A <--"交替"--> G[参数 A-G δΦ Bytes] B <--"交替"--> H[参数 A-G δΦ 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}$ 指有多少个非 PP 的 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，引入额外的开销。当然，SP 的 a2a 仅仅是 layout transform 的一种，真实通讯原语应当参考具体 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 端，跑标准的 $(s_\t{expert},H)\to(s_\t{expert},4H/N_\t{rep})\to(s_\t{expert},H)$ 操作，其中 $s_\t{expert}$ 为该 expert 收到的 token 数。
- （FP 时需要 TP all-reduce）

特别地，本节的 contract 设计和 layout transform 不一定需要手动配置。一般框架对 layout transform 的支持都很完善，还存在一些编译器可以自动搜索并行方案。

现在来玩一点花样。比如说，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 融合了。

另一个花样是把 SP(Ulysses) 和 CP(ring-attn) 融合。SP 的痛点是必须把整个 head 塞到同一张卡下，CP 的痛点是 ring 组中如果包括过多卡则效率有限。

因此有缝合二者的 head-context 2D parallelism，在 CP 的角度下就是先进行一些 a2a 通信把若干个 head 分成一组，然后每一个 head 组分配若干张 CP 卡进行 ring-attn，所有的 head 组彼此独立，内部 ring-attn 结束后再 a2a 回来。

在这个场合，$N_\t{CP}$ 和 $N_\t{SP}$ 不再能混用，二者需要进行显式区分。

III.IV. Pipelines

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

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

除此之外，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。

IV. Parallelisms at Inference

以上所有考虑的都是训练阶段的行为。现在考虑 inference。可以发现，inference 比起训练，差别是 非常大 的：

只 FP 不 BP。这意味着，激活值永远是临时的，不需要 save for BP。这也意味着 AC 和 ZeRO-R 之类技巧无用。
不需要梯度和优化器，因此 ZeRO 0/1/2 毫无价值；在 ZeRO-3 中使用的权重分片技巧倒是有一定意义，但是这也只是一种节省储存开销的可选技巧，且一般会用其它并行方法，而不是简单粗暴的分片。
仍然可以拆成多个 micro-batch 进行 PP，且因为没有 GA，不需要在 step 时把流水线清空；但存在其它仅限于 inference 阶段的 bubble 诱因。
KV Cache 成为额外需要在意的对象。

Inference 必须被拆成两个阶段：

Prefill：处理 prompt，写入 KV$，一般是 compute-heavy 的。
Decode：生成单个或一批（如果启用 MTP）token，读取并更新 KV$，同时强调 latency 和 bandwidth。

这两个阶段的具体细节将在下文阐述。

IV.O. Notations

本节统一 notation。

$R_\t{act}$：当前 active requests 数。
$S_i$：第 $i$ 个 request 本轮 forward 的输入 token 数。
$C_i$：第 $i$ 个 request 要 load 的 KV$ length，也即已处理的历史 token 数。
$S_\t{act}$：当前 active batch 的总输入 token 数。
$C_\t{act}$：当前 active batch 的总历史 KV$ token 数。
$\delta_w$：推理权重字节数，可能是 FP16/BF16/FP8/INT8/INT4。
$\delta_{kv}$：KV cache 字节数，常常不等于 $\delta_w$。
$S_{kv}$：KV cache 的每个 block 包括多少个 token。

补充服务指标：

TTFT：time to first token，主要受排队 + prefill 影响；
TPOT / ITL：time per output token / inter-token latency，主要受 decode 影响；
Throughput (吞吐)：单位时间输出 token 数；
KV capacity：在给定显存下最多同时服务多少 context token；
SLA / tail latency：调度策略是否牺牲长尾请求。

IV.I. General Frame

标准的自回归模型，不考虑 linear attn / dLLM 等特殊架构，则不管是 Prefill 还是 Decode，都可以被统一为以下框架：

每个 request 都有 $S_i$ 个输入 token 要处理。这些 token 可以是：
- 一个完整的 prefill prompt。
- decode 时，上一次 forward 新增的 autoregressive 预测。此时有 $S_i=1$（vanilla）或 $S_i$ 较小（如果用 MTP 或 speculative decoding）。
- 一个 chunked prompt。具体含义将在下文详解。
每个 request 同时在 KV$ 中有一些需要加载的 KV，长度是 $C_i$。
- 如果是 full prefill 或 chunked 的第一段，有 $C_i=0$，此时完全不需要加载 KV。
- 如果是 decode 或 chunked 的非起始段，有 $C_i>0$。
request 在 forward 的每一层都需要从 KV$ 中拉取长度为 $C_i$ 的 KV$ 并计算相应的 attn 结果。
这些 request 不会在 batch 维拼接，而是直接在 context 维拼接。逻辑上的 batch 仍然存在，但不以 batch 维度的形式出现。因此，输入形状（指过了 embedding 层后的形状）是 $(S_\t{act},H)$。
各个 request context 之间被 attn mask 保证独立；内部被 causal mask 保证因果性。PosEmb 会对每个 prompt 独立从 $C_i$ 开始计数，保证每个 prompt 的 inference 都 agnostic to 拼接。
通过合理配置 flash-attn，其开销可以写作 $\sum(S_i^2+S_iC_i)$。

虽然有这个统一框架，但具体部署时，prefill 和 decode 各自会设计特种 kernel。

产出：

首先，如果当前是 full prefill / chunked 的最后一个 chunk / decode（且未启用 Speculative Decoding），会立刻得到第一个 (vanilla) / 第一批 (MTP) 预测 token。
然后是 KV。算完后每个 request 在每一层均新增 shape 为 $(2,S_i,N_{kv},d)$ 的 KV term；如果这个 request 没有告终，则其需要被放到 KV$ 中。
最常见的模式是 paged KV$，分成 block，每个 block 的形状均为 $(2,S_{kv},N_{kv},d)$，一般取 $S_{kv}=16/32$。可简单计算得到：
- 单 token 的 KV$ 占用：$2\delta_{kv}LN_{kv}d$ Bytes。
- 整个 prefill batch 的总 block 数目：$L\sum\lceil S_i/S_{kv}\rceil$。不同 prompt 的 block 一般独立存储，方便往后面追加新 KV 项；同时如果这是 chunked 或 decode，则新 KV 项可以直接追加到已有的 block 中。
- 单 block 的存储开销：$2\delta_{kv}S_{kv}N_{kv}d$ Bytes。
和经典的 $ 逻辑相同，分为算子面对的 逻辑 block 和实际存储的 物理 block，二者之间通过 block table 建立映射。注意到这一套和 Virtual Memory 中的 page table 并没有本质不同；page table 有 cache，物理 block 也可以自动被 cache 到 GPU L1/L2$。同理，其也可以被手动 offload 到 CPU 中。

IV.II. Schedulers

在权重静态、KV 动态、请求持续进入退出的条件下，需要一些 scheduler 动态分配 GPU 运算、KV$ 存储和跨卡传输。本节将分析若干常见的 scheduler 方法。

Continuous Batching：一旦一个 decode prompt 被处理完毕，它就会被踢出当前的 batch；同理，如果有新请求到来且资源足够，则其可以直接进入 batch，下一轮的 $R_\t{act}$ 增加一。

Prefix Caching：多个 prompt 以 copy-on-write 的模式共用一个 KV$ block。这适用于 重复的前缀上下文 或 system prompt 的场景，避免重复 forward。

Chunked Prefill，在前文框架中被多次提到：

如果某个 prompt 很长（但没长到需要 CP/SP 的地步），直接把它一次性扔进 batch 会阻塞其它任务。
因此会选择把这个 prompt 分解为若干 chunk，每个 chunk 包括比如说 $512$ 个 token。一个 chunk 处理完后会直接存入 KV$，而同一个 prompt 之后的 chunk 需要从 KV$ 拉取历史 KV 信息。
可以发现，第一个之后的 chunk 的接口和 decode 一样：同样需要读取 KV$，同样需要写回；但是 $C_i$ 和 $S_i$ 不同会导致对二者的处理方式不同。
在 PD 分离的场合格外有用：因为这允许重叠 KV$ 发送与 prefill 计算。
与之相对的，不 chunk 的 prefill 被称作 full prefill。

KV management：如果当前活跃的物理 KV block 数目已经占满 GPU 的存储，则需要一些处理手段。解决方案包括：

swapping，按照某种调度策略把某些 KV$ offload 到 CPU。但是这样显然会引入额外的传输开销，且会让被 offload 的 request 的 latency 显著上升，严重伤害 tail latency。
re-computing，直接把某些请求打回 waiting list，相应的 KV$ 释放。这样显然很浪费，且同样伤害 tail latency。
限制并发度或留出冗余，确保所有被 serve 的请求都有足够资源，代价是利用率下降。

request admission：衔接 KV management 部分，动态管理什么 request 应该被优先处理的技巧。

首先有一些基础的网关技巧被迁移过来，比如说并发度、权重准入等等，没学过网关不予置评。
然后，只有在当前剩余 block 数确定能容纳完整的 prefill，且有足够置信度容纳 decode，才会被接受；如果能 prefix caching 则可以优先准入。
chunked prefill 能允许长 prompt 分段处理，和短请求一起被服务。

现在可以全面分析各个 request 类型以及其运算特点了：

Full Prefill / 第一个 Chunk：
- $C_i=0,S_i\gg 1$。
- compute-heavy，但是不需要读历史 KV$。
- 仅限本文，在以下讨论中，用类型 A 代指。
之后的 Chunks / 从 Prefix Cache 中读取：
- $C_i>0,S_i\gg1$。
- 同时包括 $S_i^2$ 的内部 attn 和 $S_iC_i$ 的 KV$ load。
- 同时读写大量 KV。
- 仅限本文，在以下讨论中，用类型 B 代指。
小 chunk 场景 / decode：
- $C_i>0$；$S_i=1$（vanilla decode）/ $S_i>1$ 但仍然较小（其它场景）。
- 包括 KV 的大量读和少量写。
- memory-bandwidth-heavy。
- 仅限本文，在以下讨论中，用类型 C 代指。

现在分析各个数据之间的联系。

latency 直接关系到 TTFT (prefill) 或 TPOT (decode)。
吞吐在 request 多时会影响排队时间，进而间接影响 TTFT。

然后是一些参数的影响。

推理的流量被拆分为以下几项：
- 权重读取，共 $$\delta_w\Phi_\text{act}$$ Bytes。
- KV 读取，共 $2\delta_{kv}LN_{kv}d C_\text{act}$ Bytes。
- 多卡通信的额外开销。
当 $S_\t{act}$ 较小时（常见于所有 request 都是 decode 的场景），流量开销少，单个请求的 TPOT 一般较好，但 GeMM 跑不满，存在闲置算力。
当 $S_\t{act}$ 较大时，吞吐上升，但单个请求的 TPOT 会下降。
进一步，当 $C_\t{act}$ 较大时，则 KV 读取可能成为主要开销，而且其显存占用可能会成为瓶颈。

虽然主流大规模并行（也是本文的主题）一般会选择把 Prefill 和 Decode 拆到不同卡上进行，称作 PD 分离；但在同一张卡同时 Prefill + Decode 也是常见选择，称作 colocated serving。

PD 分离有助于设计 domain-specific 的硬件，例如 prefill 卡相对而言更重 compute 轻 memory load，而 decode 卡的性能需求相反；但是需要额外多一步将 KV$ 跨卡传输的步骤，所以两者都需要高网络带宽。
Colocated Serving 则可以兼顾 compute 和 communication、填充流水线或者防阻塞。但代价是原本快速的 decode 可能被同一个 microbatch 中的 prefill 拖慢，TPOT 恶化，虽然吞吐可能有改善。

IV.III. Parallelisms

本节迁移前文中叙述的若干种并行方法。

DP：

因为不再需要 ZeRO 式的优化器同步，所以只需要新开一份实例，持有完整的权重副本，称作 Replica Serving。
KV$ 同样关于副本独立；特别地，prefix caching 理论上可以共享，不过要实现共享需要额外设计。
request admission 现在需要套在 replica serving 外面，动态决定请求进哪个 server。
可以启用 ZeRO-3 式的权重分片思想，但是这退化为纯粹的分布式存储技巧，会引入额外通信代价，且切分结果往往不如其它并行方法来的自然。

PP：

不需要 BP，流水线非常简单，而且激活值算完即可直接丢弃。
bubble 的主要来源不再是 FP/BP 等待、收尾的启动/清空，毕竟请求是源源不断的；取而代之的，可能来自于过少请求填不满流水线、prefill/decode 混排等。
分类型讨论：
- 首先，不管哪种类型，KV$ 都是分阶段存储的，只需要通信激活值即可。
- 类型 A：粗粒度、长耗时的任务，可能阻塞 decode。通过限制 chunksize 可以更好填充流水线。
- 类型 B：因为需要 load KV$，其计算时间会比类型 A 更加异质，为 pipeline scheduler 增加压力。如果仍然假设所有类型 B 的执行时间相同，则会导致气泡增加。
- 类型 C：通信的激活值量极小，但通信频率很高，如果不混排则难以 overlap。
- 混排：核心任务是保证每个 micro-batch 的计算时间和存储开销相近，这样才能最有效减少 bubble。把一个类型 A/B 的操作和多个类型 C 的操作混在一个 batch 里也可以平衡通信量和计算时间，更有效重叠。
同一个 request 在所有 stage 的 KV$ 逻辑是 严格同步 的，通过 全局逻辑 进行管理；不论是 eviction/offload/recovery 都是以 request 为单位进行的，不会出现在达到稳态时，一部分 stage 上该 request 已经被 evict 了，另一部分还是 active 的。
也因此，如果 stage 切分不均匀，可能发生某些层 KV$ 已经塞不下了，另外的层还绰绰有余。此时整个调度会受到容量最紧张的 stage 的限制，也即 木桶效应。
和训练时一致，最大的优势在于 p2p 通信非常适合大规模并行。进一步，因为不需要 BP，每一个 stage 的计算都可以进行极致的 分块并行，算完一小块激活值就立刻传输，进行更彻底的重叠。
可以提高稳态吞吐，但如果不考虑提升吞吐带来的等待时间减小，则 TTFT/TPOT 反倒会恶化：节点间通讯代价直接计入 prefill 时的 TTFT 和 decode 时的 TPOT，而且切分越细计入越多。

TP：

因为 TP 是 head-wise 独立的，所以 KV$ 同理，每张卡只需要维护它分到的那几个 head 的 KV$ 即可，相当于直接扩充了单层的可用 KV$ 容量。然而因为 GQA 时可能需要复制 KV head，所以扩充幅度可能不等于 TP 组内卡数。
分类讨论：
- 类型 A/B 类似，且对 TTFT 的优化明显。
- 类型 C 的 $S_\t{act}$ 一般较小，如果对 FFN 层进一步切分可能让耗时被 Kernel 启动开销、PCIe/NVLink 协议栈延迟、同步开销 吃掉，因此反倒可能恶化 TPOT。
- 混排时主要加速的仍然是 A/B，C 的优化往往不明显。
同理，同一 TP 组中的所有卡遵循 all-or-nothing，而且因为只是 head 不同，KV$ 调度更加容易。

SP(Ulysses)：

在训练的场合，它在 head 上的行为和 TP 相同，同样是 headwise 独立；但是 inference 时的行为需要定义 KV$ 的具体 contract：
- 如果选择在 QKV 的 a2a 之前储存 KV$，则其存储的 KV$ 是 sequence 分片、head 完整的，此时分析会更贴近 CP，将在下一部分详细描述。
- 否则，如果选择在 a2a 之后储存 KV$，则储存的是 sequence 完整、head 分片的 KV$，此时分析和 TP 一致。这一部分对 SP 的分析暂时只考虑这种 contract。
分类讨论：
- 类型 A/B 分析同 TP。
- 类型 C 更特殊，这种 $S_i$ 很小的场景 SP 毫无意义，因此会转而走 TP（或 Megatron-SP）。
- 进一步，混排时要动态判断所有 SP 段的长度是否足够，如果不够需要 动态切换为 TP / Megatron-SP。但这样做对调度要求极高；同时虽然对 attn 的处理相同，但 SP 和 TP 对 ffn 的处理不同；此时要么提前预制好 TP 的 ffn 分片矩阵，要么需要付出额外代价，把持有的完整 ffn 参数 reshard 为 TP 需要的部分 ffn 参数。

CP：

这是第一个行为和训练时会出现较大分歧的并行方案。在训练时，CP 切的是文本；但是在 inference 时，则不限如此。
形式化地，只考虑第 $i$ 个 request 的某个 head，则 CP 要同时考虑三个对象：
- 本轮产出的 QKV，每个 request 的长度均为 $S_i$，分布在一张或多张卡上，记作编号集合 $\c S_i$。同一个 token 的 QKV 总是同时产出的。有一个非标准术语 query parallelism 可以称呼。
- 本轮之前已经存在的历史 KV$，长度共计 $C_i$，分布在零张（如果是某个 request 的第一步）、一张或多张卡上。有一个非标准术语 cache parallelism 可以称呼。记编号集合 $\c C_i$ 为持有 最新 KV 或 历史 KV$ 的卡集合，则因为 $\c S_i$ 中的所有卡同时持有最新的 QKV，所以必有 $\c S_i\sube\c C_i$。
- 新增的长度为 $S_i$ 的 KV 要被追加入一些卡的 KV$ 上。这些卡可能就是产出 QKV 的卡，也可能是历史 KV$ 储存的卡，还可能均非。接受这些 KV 的卡的编号集合记作 $\c A_i$。
- 只要这三者有一个是分布式存储的，就可以被看做是 CP。
- 另外，本轮 attn 算完后，结果可能要被传输给另一组卡进行下一步的 FFN/MoE，这一组卡的编号集合记作 $\c F_i$。
所有 Q 和所有 KV/KV$ 都要计算得到局部的 logit。这既可以通过把 Q 发送到 KV/KV$ 端实现，也可以反过来。不同卡的结果可以使用类似 flash-attn 的手法合并，这点在传统 ring-attn 中已经见识过了。有非标准术语：发送 Q 的称作 Q-moving，发送 KV 的称作 KV-moving。
现在分类讨论。此处会列举 CP 时的若干解决方案，并对应相应的前述类型。
- $\c S_i=\c C_i=\c A_i=\c F_i$。
 - 同一组卡持有 QKV、历史 KV$，且同样是下一步 FFN 的去处。
 - 类型 A、$C_i$ 不过长的类型 B，乃至传统训练态的 ring-attn，本质上都是这种场景。自然地，此处会采用 ring-attn 作为最佳手段，选择 KV-moving，以在 GQA 等场景减少开销。但是，如果 KV$ 过大，或许不是最佳选择。
 - 虽然一般 KV$ 会就地落盘，但不排除有负载均衡的需求，此时需要 cache redistribution。
- $|\c S_i|=1,|\c C_i|>1$；$S_i$ 很小。
 - 这是 类型 C 的常见场景，request 在单张卡上，但是 KV$ 跨卡分布。此时因为 Q 少但 KV$ 多，一般会选择 Q-moving。
 - 这种风格类似 flash-decode。
- $1<|\c S_i|<|\c C_i|$。
 - 这一般发生在 类型 B 中较靠后的 chunk，此时 KV$ 已经存到了很多卡中，有很高的 cache parallelism；但是出于种种原因（比如说 chunk 不够长），query parallelism 不够高。
 - 此时有多种混合 strategy，比如说分组 Q-moving 之类，调度比较复杂。
对比：
- KV-moving：通讯开销是 $O\Big(2(S_i+C_i)N_{kv}d\Big)$，其它系数视作常数。能充分利用 GQA，更适合 $\c F_i=\c S_i$ 的场景，此时 attn 输出可以直接留在 FFN owner 端，不需要额外传输开销。
- Q-moving：通讯开销是 $O(S_iN_qd)$，其它系数视作常数。在通讯层面无法受益于 GQA，更适合 历史相对输入很长 的场景，需要 Q 的 broadcast 式传输。
- All-moving：也就是前一部分 SP 中 a2a 传输，把同一个 head 的 QKV 全部搬到同一张卡上进行。更适合 有充足 a2a 带宽 的场景。
一个值得注意的场景是，若 $|\c C_i|>|\c F_i|$，则会存在一部分卡虽然持有完整 FFN 权重，但在 FFN 阶段完全分配不到任务，是一种 结构性空转。解决方案包括通过 TP 强行扩大 $\c F_i$，再次分发 $\c F_i$ 等等，但会引起额外开销，不一定比空转更快。
进一步，出于种种原因，不同 $\c C_i$ 的处理效率可能不同，而 $\c F_i$ 端需要等待最慢的 $\c C_i$ 处理完毕才能开始 FFN，此时也会出现 等待性空转。
可能的解决方案：
- 允许不同 request 处于不同阶段，例如 request A 已经进入 FFN 但 request B 还在 attn——但需要额外的、更精细的调度，同时容易把 GeMM 拆得过细。
- 显式拆分 attn 和 FFN 组，也即 III 节提过的逻辑。
- 动态选择 CP 的并行度，不对短 context 过度切分。
- 混排，用其它任务填充 attn-only 或 FFN-only 的空缺。
- 在必要时接受空转。

EP：

训练场合的 EP 几乎总是在 DP 组上进行。但正如上述分析，DP 组已经被独立的 replica serving 替代，而跨 replica 的通信一般代价较高，因此推理场合的 EP 需要转到其它组，比如说 TP/SP/CP。
分类讨论：
- 类型 A/B：结论和训练场合基本相同。一个技巧是 热点专家复制，如果发现某个专家特别热门就把它复制多份增大它的吞吐这是训练时不好动态管理的；当然训练时的 capacity 或 drop（如果要保证输出质量，则应该禁止 drop）等技巧照样可以使用。
- 类型 C：因为太零碎了，很难 overlap，此时 EP 甚至会是对性能的负优化，这点和 TP 场合相似，甚至不如 TP：至少 TP 是负载均衡且规律的。因此，有些框架会选择在这种场合使用 TP 而不是 EP。但是，如果模型太大，出于分布式存储权重的考虑，也可能选择 EP。
- 混排：分析和 TP 相似，也即主要优化的是混排 batch 中的 prefill 请求而不是 decode。所以更倾向于 PD 分离，prefill 卡专心 EP，decode 卡不 EP、只用其它并行方案。

V. Conclusion

现在来复习一下本文的所有技术框架，制造一些 takeaways。

训练阶段的通用框架：

$N=\sum_{i=1}^{N_\t{pipeline}}N_\t{intra}^{(i)}\times N_\t{inter}$。
$N_\t{inter}$ 是非 PP 导致的同步运行的 micro-batch 的数目。
任何参数副本都需要某种形式的同步；ZeRO 是一种可选的同步方法。
如果出现副本的参数是 MoE 层，则另一种可选的同步方法是用 EP 把它拆分。EP 和 ZeRO 式同步理论上可以混排。
每一层边界处都需要显式约定具体的 contract，一般取 $\Big((b,S/N_\t{part},H)\times N_\t{part}\Big)^{N_\t{rep}}$。
朴素的 TP 不改变 contract，会有多卡持有同一份激活值的副本；进一步可以用 ZeRO-R 对激活值分布式存储，或者用 Megatron-SP 把 all-reduce 拆开，此时实际落盘的激活值可能不完全遵循 contract。
SP(Ulysses) 和 CP(ring-attn) 在 FFN 层要求的 contract 是切分输入后的 contract；前者在 attn 层内部的 contract 是切分头的，而后者仍然是切分输入的。二者也可以混用以增大 SP 的适用范围并减少 CP 通信组大小。
如果 contract 对不上，需要显式通讯进行 layout transform。

推理阶段的通用框架：

所有 request 持有 $(S_i,C_i)$，每次 pass 会为 $S_i$ 求出 KV 并落盘到 KV$，同时会产出下一轮的预测 token（如果是 full-prefill / 最后一个 chunk / decode 且未启用 speculative decoding）。
KV$ 一般以 paged 形式管理，但 evict / offload 等操作需要 request-wise 的全局管理。
DP 被 replica serving 替换。
PP 和训练时没有太大区别。
TP 和采用 head-wise 格式进行存储的 SP(Ulysses) 因为不深入 head 内部，所以不会触碰 head 内部的 KV$ 管理。
CP 要严格区分持有 Q 的卡集合 $\c S_i$、持有 KV & KV$ 的卡集合 $\c C_i$、将要落盘的卡集合 $\c A_i$ 和下一模块进行的卡集合 $\c F_i$。需要根据具体情况讨论，使用 Q-moving 或 KV-moving 之一，或者混用。
EP 一般不依托 DP 组，转而挂靠 TP/CP/SP 组。

posted @ 2026-07-03 13:08 Troverld 阅读(25) 评论(1) 收藏举报

刷新页面返回顶部

方法	并行对象	理论单卡存储 (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	/	/

Troverld

博客渲染出问题了那是博客的锅，本人什么都不会做的。

ラブパラ

O. Introduction

O.I. Notation

O.II. Collective Communication

I. Parallelisms By Themselves

I.O. Activations

I.I. Data Parallelism

I.II. Pipeline Parallelism

I.III. Tensor Parallelism

I.IV. Context/Sequence Parallelism

I.V. Expert Parallelism

II. Detailed Comparisons

II.I. Partitioning Axes

II.II. Communication Approaches

II.III. Computational Details

II.IV. Bounds & Scales

II.V. Agnostic to Implementation

III. Parallelisms With Others

III.I. Process Grid

III.II. Disentanglement Between DP and ZeRO and EP

III.III. Block Layout Contract

III.IV. Pipelines

III.V. Conclusion

IV. Parallelisms at Inference

IV.O. Notations

IV.I. General Frame

IV.II. Schedulers

IV.III. Parallelisms

V. Conclusion