第 14 讲 · ML 并行化（三）：模型并行与流水线并行 · Model and Pipeline Parallelism

对应 Maps to：CMU 15-779（Fall 2025，Zhihao Jia），Lecture 14 · Wed 10/22 · ML Parallelization: Model and Pipeline Parallelism。 必读论文：D. Narayanan et al., PipeDream: Generalized Pipeline Parallelism for DNN Training（SOSP 2019）；M. Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism（2019）。选读：D. Narayanan et al., Memory-Efficient Pipeline-Parallel DNN Training（PipeDream-2BW，ICML 2021）。 依据说明：本讲以原课 slides 13-ML-parallelization-part2.pdf 为主线：

• 数据并行的局限；模型并行（把模型切成子图放到不同设备）；
• 张量模型并行的两种切法（切输出 / 切归约维）及其与数据并行的通信代价比较；
• CNN 的并行化讨论（卷积层数据并行、全连接层张量并行）；
• Megatron-LM 切 Transformer（MLP、自注意力）和”数据 + 模型并行扩展到 512 卡”；
• 模型并行的利用率问题；GPipe 的 micro-batch 流水线、空泡公式 (p − 1)/m 与提高效率的两个方向、激活显存问题；
• 1F1B 调度；交错 1F1B 与它的空泡公式 (p − 1)/(v·m)；
• 三种并行的优缺点比较表；DeepSpeed 的 3D 并行。

slides 没有展开 PipeDream 的权重版本问题和切分算法，本讲在论文精读里补上。交错 1F1B 来自 D. Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM（SC 2021，下称 Megatron-LM 2021）；GPipe 来自 Y. Huang et al.（NeurIPS 2019）。论文数字只做转述。

parallel.py（手写的 Transformer 层及其张量并行版本、流水线调度模拟器、PipeDream 的切分动态规划）和全部实验都是本讲义自己设计的。 配套代码 Code：code/lecture14/（5 个脚本 + 共享模块 parallel.py；只需 numpy；每个脚本 < 1 秒） 预计阅读时间：3 小时（含运行代码）

0. 本讲导航 · How to read this lecture

0.1 学习目标

学完这一讲，你应该能够：

1. 说清楚 y = x·W 的两种张量并行切法（列并行、行并行）各在前向、反向通信什么，并估算它和数据并行谁的通信更少。
2. 画出 Megatron-LM 切一个 Transformer 层的方式，解释为什么每层只需要 4 次 AllReduce，以及为什么张量并行只适合放在一个节点内。
3. 推导流水线的空泡比例 (p − 1)/m，说明 GPipe、1F1B、交错 1F1B 在空泡、激活显存、通信上的取舍。
4. 解释 PipeDream 的异步流水线为什么需要权重暂存，算出每一级用的权重晚了几步。
5. 用 PipeDream 的动态规划给一个模型切分流水线级，并理解结论为什么依赖代价模型。
6. 给定模型和集群，估算 (t, p, d) 三种并行度的合理组合，并说出 Megatron-LM 2021 的三条经验。

0.2 先修与本讲范围

0.3 本讲地图

§1  为什么还要模型并行                                  ← 代码 03（演示 1）
§2  张量并行的两种基本切法                               ← 代码 01
§3  Megatron-LM：切 Transformer                          ← 代码 02
§4  流水线并行：micro-batch 与空泡（GPipe）               ← 代码 03（演示 2–3）
§5  1F1B 与交错 1F1B                                     ← 代码 03（演示 2、4–5）
§6  PipeDream：异步流水线与自动切分                       ← 代码 04
§7  三种并行的比较与 3D 并行                              ← 代码 05
§8  论文精读：PipeDream；Megatron-LM；（选读）PipeDream-2BW
§9 小结  §10 常见误区  §11 练习  §12 延伸阅读  §13 下一讲预告

0.4 关于 parallel.py

• 张量并行：手写了一个完整 Transformer 层（LayerNorm → 因果多头注意力 → 残差 → LayerNorm → GeLU MLP → 残差）的前向和反向（float64），先用有限差分检查反向是对的，再把参数切给 N 块”设备”，用第 12 讲的环形 AllReduce 实现 Megatron 的 f / g 算子，逐元素比较结果。为了简化，LayerNorm 没有可学习参数（Megatron 里它们在每块卡上复制，不涉及通信）。
• 流水线：生成 GPipe、1F1B、交错 1F1B 的调度（每台设备按什么顺序做哪些前向 / 反向），再按依赖关系模拟时间线，统计空泡、在途的 micro-batch 数（≈ 激活显存）和点对点消息数。
• PipeDream 的切分优化器：论文 §3.1 的动态规划，用穷举检查过它在小问题上是最优的。

1. 为什么还要模型并行 · Why model parallelism

slides 先回顾数据并行的问题：每块 GPU 存一份完整的模型，放不下就没法训练。第 13 讲的 ZeRO / FSDP 用”切存储”解决了放不下的问题，但它们不切计算：每块卡仍要算整个模型的前向和反向，第 3 阶段还要在网络上来回搬参数。

模型并行换一个思路：把模型切成若干个子图，放到不同的设备上，设备之间传递中间结果（激活和它们的梯度）。slides 把它分成两类：

• 张量模型并行（tensor model parallelism）：在一层之内切开参数和计算；
• 流水线模型并行（pipeline model parallelism）：在层与层之间切开，每台设备负责连续的几层。

最朴素的按层切法有一个大问题（代码 03 演示 1）：前向时只有一台设备在算，其余在等上一级的激活；反向时反过来。p = 4 时每台设备只有 25% = 1/p 的时间在工作。

设备0 |00                  aaaa                |
设备1 |  00            aaaa                    |
设备2 |    00      aaaa                        |
设备3 |      00aaaa                            |

（数字是前向，字母是反向；空格就是空闲。）§4、§5 的流水线就是为了解决这个问题。

2. 张量并行的两种基本切法 · Two ways to split a matmul

2.1 切输出与切归约维（代码 01 演示 1）

slides 用 y = x·W 说明（x：b × n，W：n × m）：

切输出（partition output，列并行）：W = [W₁ | W₂ | …]，卡 i 算 yᵢ = x·Wᵢ（y 的 1/N 列）。

• 前向：本地计算，输出是切开的；如果下一步要完整的 y，就 all-gather；
• 反向：dWᵢ = xᵀ·dyᵢ 是本地的；但 dx = Σᵢ dyᵢ·Wᵢᵀ 是各卡贡献之和，要 AllReduce（b·n 个元素）。

切归约维（reduce output，行并行）：W 按行切成 W₁、W₂……，x 相应地按列切成 x₁、x₂……，卡 i 算部分和 xᵢ·Wᵢ。

• 前向：y = Σᵢ xᵢ·Wᵢ 要 AllReduce（b·m 个元素）；
• 反向：dWᵢ = xᵢᵀ·dy、dxᵢ = dy·Wᵢᵀ 都在本地（dy 每块卡都有，dx 本来就是切开的）。

代码里两种切法和单卡的最大差分别是 3.6e-15、1.8e-15，每卡发出的元素数和公式完全一致。

关键的组合：列并行的输出是切开的，正好是行并行要的输入。把两者串起来——(x·W)·W₂，W 列切、W₂ 行切——中间结果不必拼起来，一对矩阵乘法只在末尾通信一次。这就是 Megatron-LM 切 MLP 的方式（§3）。

2.2 和数据并行比通信（代码 01 演示 2）

slides 比较了三种做法每步的通信量：

• 数据并行：同步权重梯度，∝ n·m（参数数），和 batch 无关；
• 张量并行：传激活，∝ b·m 或 b·n（token 数 × 宽度），和参数数无关。

对成对使用的 Megatron 式张量并行，每卡发出 2(N − 1)/N · b·(n + m)，数据并行是 2(N − 1)/N · n·m，所以 b < n·m/(n + m) 时张量并行更省。4096 × 4096 的投影层，这个阈值是 2,048 个 token。

slides 的结论是”最好的策略取决于模型和机器”。还要补充一点：通信量不是唯一的考虑。

• 张量并行的通信在每一层的前向、反向里都有，而且在关键路径上，很难和计算重叠；
• 数据并行的通信每步一次，可以藏进反向（第 12 讲的分桶）。

LLM 训练里每组张量并行卡每步处理的 token 数常常远大于阈值，仍然用张量并行，是为了把放不下的模型切开，而不是为了省通信。也因此它通常只放在节点内（§3.4）。

2.3 CNN：逐层选择（代码 01 演示 3）

slides 的讨论题：怎样并行 CNN？

• 卷积层：90–95% 的计算、约 5% 的参数、很大的中间激活 → 数据并行；
• 全连接层：5–10% 的计算、约 95% 的参数、很小的激活 → 张量并行。

代码按 AlexNet、VGG-16 的真实层尺寸算了一遍：

slides 的数字和 AlexNet 很接近。全部用张量并行反而是灾难（AlexNet 12 倍、VGG-16 118 倍的通信），因为卷积层的特征图太大。这个”卷积数据并行、全连接模型并行”的混合就是 Krizhevsky 2014 年的”one weird trick”——它说明每层可以有不同的并行方式，第 15 讲的 FlexFlow 就是要自动搜索这件事。

3. Megatron-LM：切 Transformer · Megatron-LM

3.1 MLP

MLP 是 Y = GeLU(X·A)，Z = Y·B。Megatron 论文先排除了一种切法：把 A 按行切、X 按列切，得到 GeLU(X₁A₁ + X₂A₂)；GeLU 是非线性的，GeLU(a + b) ≠ GeLU(a) + GeLU(b)，必须在 GeLU 之前通信一次。

它选的是：

• A 按列切：[Y₁, Y₂] = [GeLU(X·A₁), GeLU(X·A₂)]，GeLU 可以各自算，中间不通信；
• B 按行切：直接接收切开的 Y，最后 AllReduce 一次 Z = Y₁B₁ + Y₂B₂。

3.2 自注意力

多头注意力天然可以按头切：

• Q、K、V 三个投影按列切，每块卡负责一部分头，每个头的注意力完全在一块卡上算，不需要任何通信；
• 输出投影按行切，接收切开的注意力输出，最后 AllReduce 一次。

3.3 f 和 g：一层 4 次 AllReduce（代码 02 演示 1）

Megatron 把通信封装成两个”共轭”的算子：

• f：前向恒等，反向 AllReduce（放在每个并行区域的入口：输入 X 在每块卡上都有，梯度要加起来）；
• g：前向 AllReduce，反向恒等（放在出口）。

论文给出了 f 的全部实现（几行 PyTorch）：

class f(torch.autograd.Function):
    def forward(ctx, x):
        return x
    def backward(ctx, gradient):
        all_reduce(gradient)
        return gradient

一个 Transformer 层有注意力和 MLP 两个并行区域，所以前向 2 次、反向 2 次 AllReduce，每次 b·s·h 个元素。代码在 N = 2、4 上验证：

（b·s·h = 2,048；每卡发出 = 4 × 2(N − 1)/N × b·s·h。Megatron-LM 2021 把它写成每个 micro-batch 每层 8·b·s·h·(t − 1)/t。）

其他细节：

• LayerNorm、dropout、残差连接在每块卡上重复计算（输入相同），而不是一块卡算完再广播——用一点冗余计算换掉通信；
• 偏置要在 AllReduce 之后只加一次；
• 每块卡只优化自己那部分参数，不需要再同步参数。

3.4 嵌入层与并行交叉熵（代码 02 演示 2）

输出层的权重是 h × v（v 是词表大小，GPT-2 是 50,257），值得切。Megatron 按词表维切嵌入矩阵：

• 输入嵌入：每块卡只查得到自己那段词表里的 token，查不到的填 0，再 AllReduce（g）；
• 输出层：最直接的做法是 all-gather 各卡的 logits 再算交叉熵，要传 b·s·v 个数。Megatron 把输出 GEMM 和交叉熵融合：每块卡只交换每个 token 的局部最大值、指数和、目标 logit——三次 b·s 的 AllReduce。

代码里损失和单卡完全相同，通信是 all-gather 方案的 1/167（= v/6，v 越大差距越大）。

3.5 复制的激活（代码 02 演示 3）

张量并行切开了 Q、K、V、注意力概率和 FFN 的中间结果，但 LayerNorm 的输入输出和残差在每块卡上都是完整的一份。代码数了一层的激活：N = 8 时每卡的激活只降到 N = 1 的 0.27，而不是 1/8——复制的部分是一个”地板”。ZeRO 的 P_a（第 13 讲）、Megatron 后来的序列并行（Korthikanti et al., 2022：把 LayerNorm、dropout 沿序列维切开，用 reduce-scatter + all-gather 代替 AllReduce）都是为了去掉它。

3.6 什么时候划算（代码 02 演示 4）

每层的计算 ∝ b·s·h²，张量并行的通信 ∝ b·s·h，所以模型越宽，计算 / 通信比越高。代码按 Megatron 论文表 1 的四个配置估算弱扩展效率（单卡按论文的基线 39 TFLOP/s，NVSwitch 每方向 150 GB/s）：

估算偏乐观（没有计入更小 GEMM 的效率下降、kernel 启动、重复计算的 LayerNorm）。但一跨出节点（每卡 12.5 GB/s），16 路、32 路的效率就掉到 23%、12%——张量并行的通信每层都有、在关键路径上，慢网络扛不住。

slides 的”Scale to 512 GPUs by combining data and model parallelism”：Megatron-LM 用 8 路张量并行（节点内）× 64 路数据并行（节点间）训练 8.3B 的 GPT-2。

4. 流水线并行：micro-batch 与空泡 · Pipelining (GPipe)

4.1 GPipe 的想法

slides：把一个 mini-batch（每次迭代处理的样本）切成 m 个 micro-batch，让它们在 p 级流水线上依次流过，前向和反向都按 micro-batch 流水。所有 micro-batch 的梯度累积起来，在 mini-batch 结束时统一更新（流水线冲刷，flush），所以语义和单卡上的大 batch SGD 完全一样。

4.2 空泡公式

slides 的推导：设每级处理一个 micro-batch 的前向、反向时间是 t_f、t_b。

• 理想时间（每台设备的纯计算）：m·(t_f + t_b)；
• 空泡（bubble）：第一个 micro-batch 要经过 p − 1 级才到最后一台设备，最后一个反向也要经过 p − 1 级才回到第一台，空闲时间 (p − 1)·(t_f + t_b)；
• 空泡比例 = (p − 1)/m（相对理想时间）。

代码 03 演示 2（p = 4、m = 8、t_b = 2t_f）：

GPipe：空泡 0.3750（公式 0.3750）
  设备0 |0011223344556677                  hhhhggggffffeeeeddddccccbbbbaaaa          |
  设备1 |  0011223344556677            hhhhggggffffeeeeddddccccbbbbaaaa              |
  设备2 |    0011223344556677      hhhhggggffffeeeeddddccccbbbbaaaa                  |
  设备3 |      0011223344556677hhhhggggffffeeeeddddccccbbbbaaaa                      |

演示 3 扫描了 p ∈ {2, 4, 8, 16}、m ∈ {4 … 64} 的 18 种组合，模拟和公式全部一致。

4.3 降低空泡的两个方向及其代价

slides：

• 加大 m（加大 mini-batch，或减小 micro-batch）：mini-batch 太大会影响收敛，micro-batch 太小 GPU 跑不满；
• 减小 p（更浅的流水线）：每一级更大，可能放不下。

GPipe 论文的经验：m ≥ 4p 时空泡可以忽略。

4.4 激活显存

slides 指出 GPipe 的另一个问题：它先做完所有 micro-batch 的前向，每台设备要同时存 m 个 micro-batch 的激活。代码里设备 0 的在途数是 m（p = 4、m = 32 时就是 32 个）。GPipe 因此必须配合激活重计算：只存每级的输入，反向时重算。能不能改调度来省显存？

5. 1F1B 与交错 1F1B · 1F1B and the interleaved schedule

5.1 1F1B（代码 03 演示 2–3）

一前一后（one-forward-one-backward）：设备 d 先做 p − d − 1 个前向（预热），然后前向、反向交替，最后把剩下的反向做完。

1F1B：空泡 0.3750（公式 0.3750）；在途峰值 [4, 3, 2, 1]
  设备0 |00112233            aaaa44bbbb55cccc66dddd77eeee  ffff  gggg  hhhh          |
  设备1 |  001122        aaaa33bbbb44cccc55dddd66eeee77ffff  gggg  hhhh              |
  设备2 |    0011    aaaa22bbbb33cccc44dddd55eeee66ffff77gggg  hhhh                  |
  设备3 |      00aaaa11bbbb22cccc33dddd44eeee55ffff66gggg77hhhh                      |

slides 的总结：

• 稳态里一前一后，在途的 micro-batch 数被限制在流水线深度：设备 d 最多 p − d 个，和 m 无关（p = 8、m = 64 时 GPipe 64 个，1F1B 8 个）；
• 显存省了，但空泡不变：仍是 (p − 1)/m。

1F1B 来自 PipeDream；冲刷版本（PipeDream-Flush）也是 Megatron 的默认调度。

5.2 交错 1F1B（代码 03 演示 2、4）

能不能减少空泡？Megatron-LM 2021 的交错调度：每台设备不再负责连续的一段层，而是负责 v 个不相邻的块（例如 p = 4、v = 2 时，设备 0 有第 0、4 块，设备 1 有第 1、5 块……）。每个块只有一级 1/v 的计算，灌满和排空快了 v 倍：

$\text{空泡比例} = \frac{1}{v} \cdot \frac{p - 1}{m}$

代码：v = 2 时 0.1875，v = 4 时 0.0938，和公式一致。

代价（slides：“Reduce bubble time at the cost of increased communication”）：

• 每个 micro-batch 要跨 p·v − 1 个级边界而不是 p − 1 个，点对点消息从 48 增加到 112（v = 2）、240（v = 4）；
• 通信时间不可忽略时优势逐渐消失：代码里每次传输 0.4 个时间单位（前向 1）时，v = 4（40.2）比 1F1B（39.4）还慢；
• 激活显存略高：v = 2 时设备 0 在途 11 个块，相当于 5.5 级的激活（1F1B 是 4）。

Megatron-LM 2021 的摘要说交错调度能把吞吐提高 10% 以上，显存相当。

5.3 不均衡的级（代码 03 演示 5）

流水线的稳态吞吐由最慢的一级决定。代码里让 4 级中的一级慢 50%：总时间多了 39%，其他设备的忙碌比例只有 60%——而”平均每级多了多少计算”只有 12.5%。所以切分点很重要：各级计算量要尽量相等，级间要传的激活要尽量小。这引出 PipeDream 的自动切分（§6.3）。

6. PipeDream：异步流水线与自动切分 · PipeDream

6.1 不冲刷的流水线

GPipe 每个 mini-batch 结束都要冲刷，留下空泡。PipeDream（SOSP 2019）不冲刷：

• 输入级先放进 NOAM = ⌈总 worker 数 / 输入级副本数⌉ 个 mini-batch 把流水线灌满；
• 之后每台设备严格 1F1B 交替，稳态里没有空泡；
• 每个 mini-batch 的反向一做完，这一级就立刻更新自己的权重。

问题来了：同一个 mini-batch 的前向和反向之间，这一级的权重可能已经被后来的 mini-batch 更新过好几次。用旧权重算的激活配上新权重做反向，算出来的东西不是任何一组权重下的真梯度。

6.2 权重暂存与垂直同步（代码 04 演示 1–2）

• 权重暂存（weight stashing）：每一级在前向时记下用的是哪一版权重，反向时用同一版；更新照样加到最新的权重上。每一级的梯度在它用的那一版权重上是精确的，但不同级用的版本不同。
• 垂直同步（vertical sync，可选）：所有级都用”这个 mini-batch 进入第一级时”的版本，跨级一致。

论文写出了权重暂存的更新式（n 级）：

$w^{(t+1)} = w^{(t)} - \nu \cdot \nabla f\big(w_1^{(t-n+1)}, w_2^{(t-n+2)}, \ldots, w_n^{(t)}\big)$

代码在一个 4 级的小网络上真的跑了异步 1F1B（600 个 mini-batch），版本号由模拟器数出来：

和论文公式一致：第 k 级晚 n − k 步，最后一级不晚（它的前向和反向紧挨着）。

训练结果（测试损失，初始 6.84）：

如实地读这张表：

• 不暂存时梯度确实”不对”（余弦 < 1），但在这个小而平滑的问题上，学习率小时影响不大；lr = 0.5 时它发散，而权重暂存仍然收敛。
• 权重暂存在小学习率下接近顺序 SGD，学习率大时陈旧的梯度更伤（0.686 vs 0.461）。
• 垂直同步每一级都晚 3 步，比权重暂存更”旧”，在 lr = 0.2 时就发散了。论文说在他们的实验里垂直同步的影响可以忽略，默认不开；这里它反而更不稳定——延迟越大，允许的学习率越小。
• 冲刷式 1F1B 的语义最干净（大 batch 的同步 SGD），每 4 个 mini-batch 才更新一次，同样的步数下更新少，所以损失偏高。今天的大模型训练几乎都用这种同步语义。

显存：第 k 级要存 n − k + 1 份权重。论文说因为每级只有 1/n 的模型，PipeDream 的峰值显存和数据并行相当；但模型大到要靠流水线才放得下时，多存几份权重就很贵——这就是 PipeDream-2BW 的出发点（§8.3）。

6.3 自动切分（代码 04 演示 3）

PipeDream 先在单卡上剖析每层的：计算时间 T_l、输出激活大小 a_l、权重大小 w_l。然后用动态规划同时决定：

1. 层怎样切成级；
2. 每一级复制几份（级可以复制——这一级内部用数据并行，叫”1F1B-RR”，按 mini-batch 编号轮流分给副本）；
3. 放进多少个 mini-batch。

一级（层 i..j，复制 m 份）处理一个输入的等效时间（论文的式子，从 PDF 核对过）：

$T(i \to j, m) = \frac{1}{m} \max\Big(\sum_{l=i}^{j} T_l,\ \frac{2(m-1)\sum_{l=i}^{j} |w_l|}{B}\Big)$

最优流水线 = 一个用 m − m′ 块卡的最优子流水线 + 一个复制 m′ 份的最后一级：

$A(i \to j, m) = \min_{s, m'} \max\big(A(i \to s, m - m'),\ 2a_s/B,\ T(s+1 \to j, m')\big)$

论文在分层拓扑上逐层求解（先在服务器内、再跨服务器），代码实现了单层的版本，并在 30 个随机的小问题上和穷举比较，最大相对差 2.3e-16。

在 VGG-16、ResNet-50 上（16 块卡，层剖面按 V100 FP32 峰值的 50% 估算）：

（论文：VGG-16 在 Cluster-A（4 台 × 4 卡、10 Gbps）上选 “15-1”，比数据并行快 5.28 倍；Cluster-B（2 台 × 8 卡、25 Gbps）上 2.98 倍；ResNet-50 在两个集群上都选纯数据并行。）

这张表最值得学的不是具体数字，而是结论有多依赖代价模型：

• 用论文的式子（相当于 m 份副本共享一条带宽 B 的链路，比环形 AllReduce 的每卡 2(m − 1)/m 保守 m 倍），优化器自己找到了 VGG-16 的 “15-1” 结构：卷积层权重小、计算多，复制很多份；全连接层有 1 亿多个权重，只放一块卡，不必同步。
• 换成理想的环形 AllReduce，数据并行在 VGG-16 上几乎一样快，ResNet-50 总是选数据并行（和论文一致）；而论文式子下 ResNet-50 在低带宽时也会选流水线（和论文不同）。
• 论文实测数据并行的通信停顿最高占总时间的 90%，真实的 10 Gbps 以太网远达不到理想环形的效率，它的保守式子可能更贴近它的集群。

读系统论文时，要分清”结构性的结论”（权重多的层不该复制）和”依赖代价模型的数字”（快多少倍）。

7. 三种并行的比较与 3D 并行 · Comparing and combining

7.1 slides 的比较表

训练大模型需要把它们组合起来：DeepSpeed 的 3D 并行——节点内张量并行，节点间流水线并行，再在外面套一层数据并行（可以配 ZeRO-1）。

7.2 GPT-3 175B 怎么分（代码 05）

把 1024 块 A100、GPT-3 175B、batch 1536 的所有 (t, p, d) 组合（t·p·d = 1024）都用一个简单的代价模型估算：

• 计算按 Megatron-LM 2021 的 FLOPs 公式，每卡 312 TFLOP/s × 60%；
• 空泡 (p − 1)/m；
• 张量并行每层 6 次 AllReduce（含重算）；
• 流水线点对点、数据并行 AllReduce 都按不重叠算；
• 显存：16 字节 / 参数 × Ψ/(t·p) + 激活检查点。

结果：30 种组合里只有 10 种放得下。最快的几种：

（Megatron-LM 2021 实测：同样的模型、卡数和 batch，每卡 140 TFLOP/s，峰值的 45%。这里的估算偏乐观，只看排序。）

和 Megatron-LM 2021 的三条经验对照：

1. 张量并行用到一个节点的卡数（8）为止，再往上用流水线：t = 16 时 AllReduce 跨节点，张量并行一项从 2.5 s 涨到 56 s。
2. 模型并行的总规模 t·p 只要大到放得下，剩下的都给数据并行：纯数据并行每卡要 2.8 TB；在放得下的范围里（p ≥ 8），p 的选择几乎不影响时间——卡数固定时 m = B/d 随 p 成正比增长，空泡比例 (p − 1)/m ≈ N/(B·t) 几乎不变。
3. micro-batch 大小、全局 batch 都影响效率：batch 从 384 加到 3072，空泡从占计算的 31% 降到 4%，每卡从 129 升到 162 TFLOP/s——但 batch 受收敛限制。

三个维度互相牵制，没有一个旋钮能单独调到最好。可选的配置有几十种（更复杂的模型、异构的集群会有成千上万种），代价模型说了算——这正是第 15 讲”自动并行”的出发点。

8. 论文精读 · Paper reading

8.1 PipeDream: Generalized Pipeline Parallelism for DNN Training（Narayanan et al., SOSP 2019）

1. 问题：数据并行在大模型、慢网络下通信开销极高（论文实测通信停顿最高占训练时间的 90%）；传统模型并行利用率低；GPipe 式的流水线每个 mini-batch 都要冲刷。

2. 为什么困难：

• 训练是双向的流水线：前向走完还要反向走回来，前向的中间结果要留到反向；
• 朴素地流水会让同一个 mini-batch 的前向、反向用不同版本的权重；
• 频繁冲刷又浪费硬件；
• 怎样切分才能各级均衡、通信最少，取决于模型和硬件。

3. 以前的工作为什么不行：

• 数据并行通信量 ∝ 参数，大模型、慢网络下扩展差；
• 模型并行（按层切）同一时刻只有一台设备在工作；
• GPipe 每个 batch 冲刷一次，有空泡；
• 异步的数据并行（ASP、Hogwild）收敛差（论文：ASP 数据并行要比 PipeDream 多花 7.4 倍的时间才达到 48% 的精度）。

4. 新意：

• 1F1B 调度（以及级复制时的 1F1B-RR）：稳态里每台设备都在工作，在途的 mini-batch 数受限；
• 权重暂存（和可选的垂直同步）：保证每一级前向、反向用同一版权重，并写出了更新的数学形式和陈旧程度；
• 自动切分：剖析每层的计算、激活、权重，用动态规划在分层拓扑上同时决定切分、复制和在途数（论文说求解不到 8 秒）。

5. 结果：

• 在 Cluster-A（4 卡 V100 / 台、10 Gbps）、Cluster-B（8 卡 V100 / 台、NVLink、25 Gbps）上，到达目标精度的时间比数据并行快最多 5.3 倍（VGG-16，“15-1”）；
• GNMT、AWD LM、S2VT 上 1.5–4.3 倍；
• ResNet-50 上优化器选择纯数据并行（1×）——它不是总选流水线；
• 通信量比数据并行最多减少 85% 以上（VGG-16、AWD LM）。

本讲的验证：异步 1F1B 的陈旧程度和论文公式一致；动态规划在论文的代价模型下复现了 VGG-16 的 “15-1” 结构；换成理想环形 AllReduce 后结论明显变弱（§6.3）。

局限：

• 异步语义改变了优化算法（陈旧的梯度），需要逐个任务验证收敛，学习率要更保守（代码里垂直同步在 lr = 0.2 时就发散了）；
• 权重暂存要存多份权重，对超大模型不划算；
• 剖析 + 代价模型要准确，否则切分结论可能偏；
• 当时的实验都是十几到几十块卡、亿级参数的模型。

8.2 Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism（Shoeybi et al., 2019）

1. 问题：几十亿参数的 Transformer 一块卡放不下。能不能用一种简单、不需要新编译器的方法把单层切开？

2. 为什么困难：

• 层内切分要在前向、反向里插通信，插错位置要么结果不对、要么通信太多；
• 现有的模型并行框架（Mesh-TensorFlow 等）需要新的语言或编译器；
• 规模变大后训练本身也会出问题（BERT 的 LayerNorm 位置）。

3. 以前的工作为什么不行：

• GPipe 等流水线方案需要冲刷、重计算；
• Mesh-TensorFlow 需要用新的 DSL 重写模型；
• 数据并行放不下。

4. 新意：

• 利用 Transformer 的结构，只用 f / g 两个共轭算子、几行 PyTorch 代码就把 MLP、自注意力切开；
• 一层只需要前向 2 次、反向 2 次 AllReduce；
• 输入 / 输出嵌入按词表切，输出 GEMM 与交叉熵融合，只传标量；
• LayerNorm、dropout、残差在各卡上重复计算以省通信；
• 张量并行与流水线并行正交、可以组合。

5. 结果：

• 单卡基线 39 TFLOP/s（峰值的 30%）；
• 8 路张量并行的 8.3B 模型达到线性扩展的 77%；加上 64 路数据并行、共 512 块 V100 时为 74%（摘要写作 76%），整体 15.1 PFLOP/s；
• 8.3B 的 GPT-2 在 WikiText103 上困惑度 10.8（之前最好 15.8）、LAMBADA 准确率 66.5%（之前 63.2%）；
• 3.9B 的 BERT 在 RACE 上 90.9%；
• 发现 BERT 变大时要调整 LayerNorm 的位置才能继续变好。

本讲的验证：手写的 Transformer 层在 N = 2、4 上切分后和单卡一致到 1e-15，每层恰好 4 次 AllReduce；并行交叉熵只传 v/6 分之一的数据；估算的效率随 h 增长、跨节点时崩溃。

局限：

• 张量并行的通信每层都有、在关键路径上，只适合节点内的高带宽互连，扩展度有限（一般 ≤ 8）；
• 复制的 LayerNorm、残差让激活显存有一个地板（后来的序列并行解决）；
• 要求头数、FFN 宽度能被卡数整除。

8.3 两篇论文之间的联系

• 两个正交的切法：Megatron 在层内切（张量并行），PipeDream 在层间切（流水线并行）。前者通信频繁但没有空泡，后者通信少但有空泡（或陈旧）。Megatron 论文自己就说两者”正交且互补”，后来的 Megatron-LM 2021 把它们和数据并行组合成 PTD-P（§7）。
• 两种对”同步”的态度：Megatron 坚持完全同步（结果和单卡一致）；PipeDream 接受有界的陈旧换取没有冲刷。历史的选择是同步的那一边——今天的大模型训练用的是 PipeDream-Flush（1F1B + 冲刷）加交错调度，而不是异步的 PipeDream。原因之一：大模型的训练太贵，没人愿意在收敛性上冒险。
• 都在用”结构”省通信：Megatron 利用 GeLU 是逐元素的、注意力头之间独立；PipeDream 利用”卷积层权重少、全连接层权重多”来决定复制谁。一个可以讨论的问题：PipeDream 的切分优化器能不能把张量并行也纳入搜索空间？这正是第 15 讲 Alpa 做的事。

8.4 选读：Memory-Efficient Pipeline-Parallel DNN Training（PipeDream-2BW, ICML 2021）

• 问题：PipeDream 的权重暂存要存最多 p 份权重，对十亿级参数的模型太贵；GPipe 的冲刷又浪费吞吐。
• 方法：双缓冲权重（2BW）——只保留两个版本，每 m 个 micro-batch 合并梯度、更新一次；新进入流水线的 micro-batch 用新版本，还在途的继续用旧版本。更新的语义是固定延迟 1 步：w^(t+1) = w^(t) − ν·∇f(w^(t−1))。还提出了 PipeDream-Flush（1F1B + 冲刷，同步语义、显存更低）作为另一个选项，以及一个考虑显存约束的切分规划器。
• 结果：训练大型 GPT、BERT 最多快 20 倍，最终精度相近。
• 意义：它介于 GPipe（同步、有空泡）和 PipeDream（异步、多份权重）之间，也说明”异步流水线”这条路在大模型上的代价主要在显存。

9. 小结 · Summary

1. 模型并行切计算：张量并行在层内切，流水线并行在层间切。朴素地按层切，p 台设备的利用率只有 1/p。
2. 张量并行的两种基本切法：列并行（前向不通信、反向 AllReduce 输入梯度）、行并行（前向 AllReduce、反向不通信）；成对使用时一对矩阵乘法只通信一次。通信 ∝ 激活，数据并行 ∝ 参数，各有适用的场景；CNN 的卷积层用数据并行、全连接层用张量并行。
3. Megatron-LM：注意力按头切、MLP 列切 + 行切，f / g 算子每层 4 次 AllReduce，结果和单卡一致；词表并行的交叉熵只传标量。张量并行适合节点内（≤ 8 路），跨节点效率崩溃。
4. 流水线的空泡 (p − 1)/m：GPipe 和 1F1B 一样；1F1B 把在途的 micro-batch 从 m 降到 p；交错 1F1B 把空泡降到 (p − 1)/(v·m)，代价是 v 倍的点对点通信。最慢的一级决定吞吐。
5. PipeDream：不冲刷的异步 1F1B 需要权重暂存，第 k 级用的权重晚 n − k 步；动态规划同时决定切分和复制，但结论高度依赖代价模型。
6. 3D 并行：张量并行放节点内，流水线跨节点，模型并行刚好放得下以后剩下的给数据并行。几十种配置互相牵制，引出自动并行。

10. 常见误区 · Common misconceptions

• “张量并行会改变计算结果。” 不会。代码里输出、所有梯度和单卡一致到 1e-15，它只改变了谁算哪一部分。
• “张量并行越多越好。” 每层都要 AllReduce，而且在关键路径上；一跨出节点效率就从 88% 掉到 23%（代码 02 的估算）。
• “1F1B 比 GPipe 的空泡小。” 一样大，都是 (p − 1)/m。1F1B 省的是激活显存（在途的 micro-batch 从 m 降到 p）。
• “交错调度总是更快。” 它用 v 倍的点对点通信换空泡；通信慢的时候反而更慢（代码 03 演示 4）。
• “流水线并行没有通信。” 有，只是少：每个 micro-batch 在每个级边界传一次激活（前向）和梯度（反向），量是 b·s·h，远小于数据并行的参数同步。
• “PipeDream 的异步流水线和数据并行等价。” 不等价：权重暂存下第 k 级的梯度晚 n − k 步，是带陈旧的 SGD，需要单独验证收敛、可能要更小的学习率。
• “切分优化器给出的加速比是可靠的。” 它只在代价模型准确时可靠。同一个 VGG-16，论文的式子说快 17 倍，理想环形 AllReduce 说只快 7%。

11. 练习 · Exercises（含参考答案）

练习 1　y = x·W，x：b × n，W：n × m，N 块卡。分别写出列并行、行并行在前向和反向各需要什么通信、每卡发出多少元素（环形算法）。

答案：

• 列并行：前向本地（要完整 y 时 all-gather，(N − 1)/N · b·m）；反向 dx 要 AllReduce，2(N − 1)/N · b·n。
• 行并行：前向 y 要 AllReduce，2(N − 1)/N · b·m；反向本地。
• 代码 01 演示 1 里（b = 8、n = 12、m = 16、N = 4）分别是 96、144、192 个元素。

练习 2　为什么 Megatron 的 MLP 第一层按列切、第二层按行切，而不是反过来？

答案：第一层按列切，每块卡得到 GeLU 输入的一部分列，GeLU 是逐元素的，可以各自算，中间不需要通信；第二层按行切，正好接收切开的输入，最后 AllReduce 一次。反过来（第一层按行切）会得到部分和 X₁A₁ + X₂A₂，GeLU 不是线性的，必须先 AllReduce 再算 GeLU，多一次通信。

练习 3　一个 Transformer 层 b = 8、s = 2048、h = 8192，8 路张量并行，BF16。每个 micro-batch 每层每卡要发出多少字节？在每方向 300 GB/s 的 NVLink 上大约多久（忽略 α）？

答案：

• 4 次 AllReduce × 2 × 7/8 × b·s·h = 7 × 8 × 2048 × 8192 ≈ 9.4 × 10⁸ 个元素 ≈ 1.88 GB；
• 1.88 GB / 300 GB/s ≈ 6.3 ms；
• 每层每卡的计算约 3 × 24 × b·s·h² / 8 ≈ 9.9 × 10¹² FLOPs，在 200 TFLOP/s 下约 50 ms，通信约占 11%。跨节点（25 GB/s）时通信变成 75 ms，比计算还长。

练习 4　p = 8 级流水线，希望空泡比例低于 10%。1F1B 至少要多少个 micro-batch？交错 v = 4 呢？两者设备 0 各要存多少个 micro-batch 的激活？

答案：

• 1F1B：(p − 1)/m < 0.1 → m > 70，至少 71（交错调度还要求 m 是 p 的倍数）；设备 0 存 p = 8 个。
• 交错 v = 4：(p − 1)/(4m) < 0.1 → m > 17.5，又要是 8 的倍数，取 24；设备 0 存的块数会多一些（代码里 p = 4、v = 2 时是 11 个块 ≈ 5.5 级）。
• GPipe 要存 m 个，71 个 micro-batch 的激活通常放不下，只能重计算。

练习 5　PipeDream 的 4 级流水线（权重暂存），第 2 级在第 100 个 mini-batch 的前向里用的是第几版权重？如果开了垂直同步呢？

答案：权重暂存：第 k 级晚 n − k 步，第 2 级晚 2 步，用”更新 98 次后”的权重。垂直同步：所有级都用这个 mini-batch 进入第 1 级时第 1 级的版本，第 1 级晚 3 步，所以是”更新 97 次后”的权重。代码 04 演示 2 的表格可以直接验证。

练习 6　用 PipeDream 的式子，一级复制 m 份、权重 W 字节、计算 C 秒、带宽 B。m 多大时这一级从”计算受限”变成”通信受限”？换成环形 AllReduce 的式子呢？

答案：

• 论文式子：2(m − 1)W/B > C → m > 1 + CB/(2W)；
• 环形：2(m − 1)/m · W/B > C → 只要 2W/B > C（m 大时），否则永远计算受限。
• 论文的式子里通信随 m 线性增长，复制越多越贵；环形式子里每卡的通信量几乎不随 m 增长。所以用论文的式子会更倾向于少复制、多切级。

练习 7　为什么 3D 并行里张量并行放在节点内、流水线并行放在节点间，而不是反过来？

答案：张量并行每层前向、反向各 2 次 AllReduce，每次 b·s·h，而且在关键路径上，需要 NVLink 这样的高带宽；流水线每个 micro-batch 在每个级边界只传一次 b·s·h（前向）和一次（反向），量少、可以和其他 micro-batch 的计算重叠，节点间的网卡就够了。反过来，张量并行跨节点会让每层都等慢网络（代码 05：t = 16 时每步从 28 s 变成 81 s）。

练习 8　（开放）PipeDream 的异步语义在 2019 年的 VGG、GNMT 上没有影响收敛，但今天的大模型训练几乎都用同步的 1F1B + 冲刷。你认为有哪些原因？

参考思路：

• 大模型训练一次要花几百万美元，收敛性上的任何风险都不可接受，也没法逐个任务调学习率；
• 权重暂存要存多份权重，大模型上显存代价太高（2BW 就是为了这个）；
• 交错调度、更大的 batch（更多 micro-batch）已经把冲刷的空泡压到很小；
• 大 batch 训练的优化技巧（学习率预热、Adam / LAMB）都是在同步语义下发展的；
• 可复现性和调试：同步训练的结果是确定的，便于定位问题。

12. 延伸阅读 · Further reading

• Y. Huang et al., GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism（NeurIPS 2019）——micro-batch 流水线与重计算。
• D. Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM（SC 2021）——交错 1F1B、PTD-P、三条经验，3072 块 A100 上 1T 参数 502 PFLOP/s。
• D. Narayanan et al., Memory-Efficient Pipeline-Parallel DNN Training（ICML 2021）——PipeDream-2BW、PipeDream-Flush。
• V. Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models（MLSys 2023）——序列并行与选择性重计算。
• A. Krizhevsky, One weird trick for parallelizing convolutional neural networks（2014）——卷积数据并行、全连接模型并行。
• N. Shazeer et al., Mesh-TensorFlow: Deep Learning for Supercomputers（NeurIPS 2018）——用”网格 + 维度映射”统一描述张量切分，Alpa、GSPMD 的前身。
• P. Qi et al., Zero Bubble Pipeline Parallelism（ICLR 2024）——把反向拆成”输入梯度”和”权重梯度”两部分，进一步消除空泡。

13. 下一讲预告 · Next lecture

这一讲手工分析了几种并行方式，又在 05 号脚本里枚举了 30 种 (t, p, d) 组合。但真实的模型不是一摞相同的 Transformer 层：有嵌入层、有专家、有不同宽度的卷积，每个算子都可以有自己的切法，集群也可能是异构的。第 15 讲《自动并行》介绍两种把这件事交给机器的方法：

• FlexFlow：定义 SOAP（样本、算子、属性、参数）四个维度的搜索空间，用执行模拟器快速评估一个策略，用 MCMC 搜索；
• Alpa：把并行分成”算子间”（流水线）和”算子内”（张量 / 数据并行）两层，上层用动态规划切流水线级，下层用整数线性规划选每个算子的切法。

我们会在一个小计算图上实现执行模拟器和 MCMC 搜索，再写一个两级搜索的 Alpa 缩小版，看它们能不能自己”发现”本讲的手工策略。