第 17 讲 · 分布式训练 · Distributed Training

对应 Maps to：CMU 10-423 / 10-623 / 10-723 Generative AI（Matt Gormley、Aran Nayebi；slides 取自 Henry Chai 和 Pat Virtue），Lecture 17 · Wed 18-Mar · Distributed training（slides：lecture17-distributed）。 阅读材料：Shoeybi et al. (2019), Megatron-LM；Huang et al. (2019), GPipe；Narayanan et al. (2019), PipeDream；Narayanan et al. (2021), Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM；Rajbhandari et al. (2020), ZeRO；Li et al. (2021), TeraPipe。 依据说明：本讲按原课第 17 讲公开的 slides 展开：训练 Llama 的成本回顾（Meta 的 24,576 块 H100 集群）；A100 / H100 / B200 的对比与”GPU 之间的通信是主要瓶颈”；LLM 训练里的几种并行；数据并行及其根本问题（Llama-4 Maverick / Behemoth 的显存）；模型（张量）并行：Transformer 层的结构、MLP 块的列并行与行并行（并行矩阵乘法、前向与反向、两个观察、f 和 g）、注意力块按头切分；流水线并行（朴素实现的空闲、微批次、1F1B、权重暂存、冲刷、怎样切层）；“数据并行仍然是王者”；优化器并行（Adam、混合精度、K = 12、ZeRO-DP）；Llama-4 的规模表；token 并行（TeraPipe）；专家并行（容量因子）。所有实验是本讲义自己设计的，不代表课堂上讲了什么。 配套代码 Code：code/lecture17/（5 个脚本 + 1 个模块；都只要几秒） 预计阅读时间：2.5 小时（含运行代码）

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

0.1 学习目标 Learning goals

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

1. 算出一个模型训练时要多少显存（权重、梯度、优化器状态、激活值），说出为什么必须分布式。
2. 描述数据并行、张量并行、流水线并行、ZeRO 各自切的是什么、要通信什么、通信多少。
3. 用 slides 的公式推出 Megatron 的列并行 / 行并行 MLP 的前向和反向，说明 f 和 g 两个算子。
4. 画出 GPipe 和 1F1B 的时间表，算出气泡比例 (p−1)/(m+p−1)，解释 1F1B 省了什么、异步流水线为什么要暂存权重。
5. 写出 ZeRO 各阶段每块 GPU 的显存 $(2+2+K)\Psi$ 等公式；说明为什么 P_os、P_os+g 不增加通信、P_os+g+p 多 50%。
6. 说出几种并行怎样组合（张量并行在节点内、流水线跨节点、其余给数据并行）。

0.2 先修与本讲范围 Prerequisites & scope

0.3 本讲地图 Map

§1  为什么要分布式：显存账和通信账                                 ← 代码 01
§2  数据并行                                                        ← 代码 02
§3  模型（张量）并行：MLP 和注意力                                   ← 代码 02
§4  流水线并行：气泡、1F1B、权重暂存、切层、token 并行               ← 代码 03
§5  优化器并行：ZeRO                                                 ← 代码 04
§6  专家并行与组合（3D 并行）                                        ← 代码 05
§7–11  小结、误区、练习、延伸阅读、下一讲

本讲的代码不需要真的多块 GPU：dist_lab.py 里”第 i 块 GPU”就是列表里的第 i 个元素，GPU 之间的通信由 Comm 执行（结果用 torch 精确地算）并按环形算法记账。于是每一种并行都能真实地算一遍，再和单卡的结果比较。

1. 为什么要分布式 · Why distribute

1.1 成本与硬件

回顾第 15 讲：训练 Llama 要几百万 GPU 小时。slides 引用 Meta 的公告：“新的两个集群各有 24,576 块 Nvidia H100——比原来的 16,000 块 A100 的集群大得多”。slides 的硬件表：

（slides 的算力是带稀疏加速的数字；稠密 BF16 是一半左右：A100 312、H100 989、B200 约 2250 TFLOP/s。slides 里 H100 的显存带宽写的是 1.6 TB/s，和 A100 一样；H100 SXM 的规格是 3.35 TB/s。代码 01 用的是稠密算力和 3.35 TB/s。）

slides 的关键结论：GPU 之间的通信是分布式系统的主要瓶颈（加速它很值钱）。 H100 每传 1 个字节的时间里能做约 1100 次 BF16 运算（代码 01 第 3 部分）。

1.2 显存账（代码 01）

训练时每个参数要存 16 字节（混合精度 Adam：BF16 权重 2 + BF16 梯度 2 + FP32 主权重 4 + 一阶矩 4 + 二阶矩 4；P1）：

slides：“Maverick：每个参数 2 字节……4000 亿参数要 8000 亿字节 ≈ 800 GB！而我们的 GPU 只有 192 GB 显存。办法：把模型切开放在多块 GPU 上！”

激活值也很大：按 Korthikanti et al. 的估算，GPT-3 每层存 $sbh(34 + 5as/h)$ 字节（s = 序列长度、b = 批大小、h = 隐藏维度、a = 头数），一个 2048 token 的序列每层 2.87 GB、96 层 275 GB（不做激活重计算时）。所以实践中都要激活重计算（只存每层的输入，反向时重算），并把层和序列切到多块 GPU 上。

1.3 通信账

一个例子（代码 01）：7.5B 参数的模型、64 块 H100 做数据并行、每块每步 16,384 个 token。一步的计算约 1.86 秒（40% MFU）；梯度（BF16，15 GB）的环形 all-reduce 在 NVLink 上 0.07 秒（4%），在跨节点的 InfiniBand 上 0.59 秒（32%）。数据并行每步只通信一次，还算可以接受；带宽一低、或者通信更频繁（张量并行每层都通信），通信就成了瓶颈。

目标（slides）：把训练 LLM 的工作分给多块 GPU，同时让 GPU 之间的通信尽量少。好消息：基于 Transformer 的 LLM 很容易并行。容易利用的并行有：数据并行、模型（张量）并行、流水线并行、基于优化器的并行、token 并行、专家并行。

2. 数据并行 · Data parallelism

做法（slides）：训练时把每个小批次平均分给几块 GPU，每块 GPU 对自己的那些数据点算前向和反向；然后把梯度求和、广播（all-reduce）。

• 最简单、最有效的并行（只要放得下）；
• 每次迭代只有一次 GPU 之间的通信。

为什么等价：损失是对样本的平均，$\nabla\frac{1}{B}\sum_{i=1}^{B}\ell_i = \frac{1}{N}\sum_{k=1}^{N}\Big(\frac{N}{B}\sum_{i\in \text{第 }k\text{ 块}}\nabla\ell_i\Big)$——每块 GPU 算自己那 1/N 批的平均梯度，再对 N 块求平均，就是整批的梯度。代码 02 第 1 部分：4 块 GPU、梯度 all-reduce 后求平均，和整批的梯度最大相差 3.0e-08（float32 加法顺序不同的舍入）；每块 GPU 发出 $2(N-1)/N \times$ 参数字节。

根本问题（slides 的问题）：每块 GPU 都要存一份完整的模型。Maverick 的权重 800 GB，一块 192 GB 的 GPU 根本放不下——要把模型本身切开。

3. 模型（张量）并行 · Model (tensor) parallelism

3.1 切哪里

做法（slides）：对一批数据，把一层之内的前向和反向计算分到几块 GPU 上；也叫张量并行。Transformer LM 的每一层有几个子层：注意力、前馈网络、层归一化、残差连接；其中两个主要模块可以并行：MLP 块和注意力块。

3.2 并行的矩阵乘法

slides 的”旁注”：$C = AB$，$C_{i,j} = \sum_k A_{i,k}B_{k,j}$。

• 按列切 B：$B = [B_1, B_2]$，$C = [AB_1, AB_2]$——每块 GPU 算 C 的一部分列，拼起来；
• 按行切 B（同时把 A 按列切）：$A = [A_1, A_2]$、$B = \begin{bmatrix}B_1\\B_2\end{bmatrix}$，$C = A_1B_1 + A_2B_2$——每块 GPU 算一个完整形状的部分和，加起来。

3.3 MLP 块

slides：给定一批训练数据 $X \in \mathbb R^{b\times k}$ 和权重 $W \in \mathbb R^{k\times n}$，要并行的操作是前向 $Y = XW$ 和反向 $\partial\ell/\partial X = GW^\top$（$G = \partial\ell/\partial Y$）。两种切法：

列并行 $W = [W_1, W_2, \ldots, W_N]$：

• 前向：$XW = [XW_1, \ldots, XW_N]$——各算各的，不通信；
• 反向：$GW^\top = [G_1, \ldots, G_N]\begin{bmatrix}W_1^\top\\ \vdots\\ W_N^\top\end{bmatrix} = \sum_{i} G_iW_i^\top$——要求和。

行并行 $W = \begin{bmatrix}W_1\\ \vdots\\ W_N\end{bmatrix}$（X 也按列切成 $[X_1, \ldots, X_N]$）：

• 前向：$XW = \sum_i X_iW_i$——要求和；
• 反向：$GW^\top = [GW_1^\top, \ldots, GW_N^\top]$——各算各的，不通信。

两个关键观察（slides）：(1) 列并行和行并行的前向、反向计算一模一样，只是反过来；(2) 相邻的两个 MLP 层可以交替使用列并行和行并行，尽量减少跨 GPU 的求和。

Megatron 的 MLP $Z = \mathrm{GeLU}(XA)\,B$：第一层 A 按列切（每块 GPU 得到 $Y_i = \mathrm{GeLU}(XA_i)$——GeLU 是逐元素的，可以各算各的，这正是先列后行的原因），第二层 B 按行切（$Z = \sum_i Y_iB_i$）。整个 MLP 只要两个通信算子：

• f：前向是恒等（每块 GPU 拿到同一个 X），反向是求和（all-reduce $\partial\ell/\partial X$）；
• g：前向是求和（all-reduce Z），反向是恒等。

代码 02 第 2 部分按上面的公式手写了前向和反向：单卡的手写反向和 PyTorch 自动求导差 1.4e-14；张量并行切成 2 块、4 块，Z、∂X、∂A、∂B 和单卡差 2.8e-14（float64）；一个 MLP 前向、反向各一次 all-reduce；每块 GPU 只存 1/t 的权重（4 块时 512 个，单卡 2048 个）。

3.4 注意力块

slides：多头注意力天然可以按头切分——第 i 块 GPU 拿 $W_Q, W_K, W_V$ 的第 i 组列，独立算自己那几个头；各头的输出是横着拼起来的，所以可以直接接一个行并行的输出投影 $W_O$。代码 02 第 3 部分：8 个头切成 2 / 4 / 8 块，和单卡差 ≤ 6.7e-16，前向只 all-reduce 一次。

所以一个 Transformer 层：前向 2 次 all-reduce（注意力的输出投影、MLP 的第二层），反向 2 次。

3.5 张量并行的通信有多贵

代码 02 第 4 部分：GPT-3 规模的一层（h = 12288、序列 2048、批 1、BF16）切成 8 块。每次 all-reduce 的张量是 $s\cdot b\cdot h$ 个数，每块 GPU 每层发出 352 MB：

张量并行每层都通信，而且在关键路径上（要等它传完才能接着算下一个子层），所以只在一个节点里（NVLink 连着的 8 块 GPU）用。

4. 流水线并行 · Pipeline parallelism

4.1 气泡

slides：上面的张量并行其实也隐含着流水线并行——不同的层或模块放在不同的 GPU 上。问题：层与层之间的计算本来就是串行的，不是并行的！朴素的实现有大量空闲的 GPU 时间。

办法：同时处理多个微批次（microbatch）——GPipe（Huang et al., 2019）：把一批切成 m 个微批次，像流水线一样依次流过各个阶段，先做完全部前向、再做全部反向。代码 03 第 1 部分（4 个阶段，前向 1 个单位、反向 2 个单位；数字 = 微批次的前向，字母 = 它的反向，· = 空闲）：

朴素（1 个微批次）：气泡 75%              GPipe（4 个微批次）：气泡 42.9%
GPU 1  0·········aa                       GPU 1  0123·········aabbccdd
GPU 2  ·0······aa··                       GPU 2  ·0123······aabbccdd··
GPU 3  ··0···aa····                       GPU 3  ··0123···aabbccdd····
GPU 4  ···0aa······                       GPU 4  ···0123aabbccdd······

1F1B（8 个微批次）：气泡 27.3%，每段同时存着的激活 4 / 3 / 2 / 1 份
GPU 1  0123······aa4bb5cc6dd7ee·ff·gg·hh
GPU 2  ·012····aa3bb4cc5dd6ee7ff·gg·hh··
GPU 3  ··01··aa2bb3cc4dd5ee6ff7gg·hh····
GPU 4  ···0aa1bb2cc3dd4ee5ff6gg7hh······

气泡比例：一共 m 个微批次、p 个阶段，每个阶段真正干活的时间是 $m(t_f + t_b)$，总时间是 $(m+p-1)(t_f+t_b)$（最后一个微批次要比第一个晚 p − 1 个”格子”进入最后一段，反向同理），所以

4.2 1F1B

slides：一个前向、一个反向（1F1B）：稳定以后，每块 GPU 交替做一个前向和一个反向。它的气泡和 GPipe 一样，但每个阶段同时在路上的微批次最多 p 个（GPipe 是 m 个），激活值的显存小得多——于是可以用更多的微批次，把气泡压下去。slides 引用 Narayanan et al. (2021)：即使前向和反向的计算量不同（通常反向是前向的 2 倍），1F1B 仍然很高效；每个小批次结束时有一次冲刷（flush）：等所有微批次做完，再统一更新参数，下一个小批次的前向用新的参数。

代码 03 第 2 部分真的把一个 8 层的 MLP 切成 4 段、用 1 / 4 / 8 个微批次跑一遍：和整批训练的梯度最多差 4.5e-08。流水线只改变了计算的先后，没有改变算什么（1F1B 和 GPipe 也只是同一组加法换了顺序）。

4.3 权重暂存（PipeDream）

slides：如果每做一次反向就更新权重（不冲刷、异步的流水线，PipeDream），一个微批次前向用的权重和反向用的权重可能不一样；越靠前的 GPU 越严重；可能导致收敛变差。办法：权重暂存（weight stashing）——每次前向之后把用到的权重存一份，对应的反向时再拿出来用。

代码 03 第 3 部分：在 1F1B 的稳定阶段，一个微批次的前向和反向之间，各段的权重被更新了几次：

第一段的版本差 p − 1、最后一段 0，正是 slides 说的”越靠前越严重”。

4.4 怎样把层分给 GPU

slides 的另一个问题：怎样把层分到各块 GPU 上？ 办法：先测出每层的耗时，再用动态规划。流水线的速度由最慢的一段决定，所以要让最慢的一段尽量快：

（前 i 层切成 k 段时最慢一段的最小值。）代码 03 第 4 部分：12 层，嵌入层（3）和输出层（4）特别慢、其余每层 1，分成 4 段：按层数平均切（每段 3 层）时每段耗时 5 / 3 / 3 / 6，最慢 6；动态规划切成 0–1、2–5、6–9、10–11，每段 4 / 4 / 4 / 5，最慢 5。

4.5 到目前为止

slides（取自 Yuanzhi Li）：归根结底，数据并行仍然是王者；张量并行和流水线并行是用来支撑”把更多数据推过前向和反向”的——它们让大模型放得下，剩下的 GPU 全用来做数据并行。

5. 优化器并行：ZeRO · Optimizer parallelism

5.1 Adam 要存什么

slides：如果只需要计算和存储参数、梯度，数据 + 模型 + 流水线并行（通常）就够了——比如小批量 SGD 的更新只要梯度和当前的权重：$\theta^{(t+1)} \leftarrow \theta^{(t)} - \gamma\nabla\ell(\theta^{(t)})$。但很多高级的优化算法要存额外的状态。Adam 结合了 SGD 的动量（记住之前的梯度）和 RMSProp（按之前的梯度调整步长）：

$m_{-1}, v_{-1}$ 初始化为全 0。混合精度训练时，$m_t, v_t$ 和 $\theta$ 一样大，通常用 FP32 存（因为 $v_t$ 里有平方项，FP16 的范围不够）；再加一份 FP32 的主权重——K = 12 = 3 × 4 字节 / 参数（slides）。加上 FP16 的参数和梯度各 2 字节，每个参数 $(2 + 2 + K) = 16$ 字节，优化器状态占了 3/4。

5.2 ZeRO-DP

ZeRO（Zero Redundancy Optimizer powered Data Parallelism；Rajbhandari et al., 2020）的观察：数据并行的 N_d 块 GPU 各存一份完全一样的参数、梯度、优化器状态，是巨大的冗余。三个阶段，逐步把它们切开：

为什么 P_os、P_os+g 不增加通信：all-reduce 本来就等于 reduce-scatter + all-gather（P4）。ZeRO 把这两半拆开用：先 reduce-scatter 梯度（第 i 块 GPU 只拿到第 i 段梯度的总和），用它更新自己那一段的参数（只有它有这一段的优化器状态），再 all-gather 更新后的参数——总量和一次 all-reduce 一样。P_os+g+p 的参数也是切开存的，前向、反向用到某一层时才 all-gather 它的参数、用完就扔，多了一次 all-gather：3Ψ，多 50%。

代码 04 第 1 部分真的实现了这四种做法（4 块”GPU”、一个 1633 个参数的 MLP、5 步 Adam）：

（每个参数的字节数比 16 / 7 / 5.5 / 4 多一点点，是因为参数个数补到了 4 的倍数。）ZeRO 没有改变算法，只改变了”谁存哪一段、谁更新哪一段”。

论文图 1 的例子（Ψ = 75 亿、K = 12，每块 GPU，不含激活值；代码 04 第 2 部分）：

P_os 一步就省下最多（优化器状态是大头）；P_os+g+p 的显存和 GPU 数成反比——GPU 足够多时，任意大的模型都放得下，代价是更多的通信。PyTorch 的 FSDP 就是 P_os+g+p 的一种实现。

6. 更多的并行，以及怎样组合 · More parallelism

6.1 Llama-4 的规模

slides 的表最后一行：LLaMA-4（Meta，2025），30 万亿 token、2 万亿参数（2880 亿激活参数）——这是一个 MoE。

6.2 token 并行（TeraPipe）

slides：洞察：在因果注意力的 Transformer 里，某一层里每个 token 的计算只依赖这一层之前的 token。想法：不必等上一层整层算完，第 l − 1 层的 token 1 … t − 1 一算完，就可以开始算第 l 层的 token t。直觉：把流水线的粒度切得更细，空闲就更少；序列越长收益越大，而模型的上下文越来越长。（slides 的图用的是”全部前向、再全部反向”的调度，不需要权重暂存，但空闲多一些。）

代码 03 第 5 部分（近似：每块的计算按 1/c 算，忽略后面的 token 块要看前面所有 token 的额外注意力）：2 个序列、4 个阶段，每个序列切成 1 / 2 / 4 / 8 块，气泡 60.0% / 42.9% / 27.3% / 15.8%——相当于把微批次数从 2 变成了 2c。

6.3 专家并行

slides：把每个专家放在不同的 GPU 上（第 16 讲）。关键：为了让各 GPU 的负载均衡，每个专家有一个容量；分给它的 token 超过容量时，多出来的 token 直接经过残差连接传到下一层。容量因子太小，容易超额分配；太大，一些 GPU 会闲着。第 16 讲代码 04 量过：每个 MoE 层两次 all-to-all，Mixtral 规模的一层每块 GPU 发出约 235 MB，跨节点时通信 / 计算 ≈ 0.32。

6.4 组合起来：3D 并行（代码 05）

真实的大模型训练（Megatron-LM，Narayanan et al., 2021）同时用三种：张量并行 t × 流水线并行 p × 数据并行 d = GPU 总数（再加 ZeRO-1 切优化器状态）。代码 05 用一个粗略的模型，在 1024 块 H100（每台机器 8 块）上、全局批 1536 个 2048 token 的序列，枚举 (t, p, d)：

• 每块 GPU 的显存 = BF16 参数和梯度 $4\Psi/(tp)$ + 优化器状态 $12\Psi/(tpd)$ + 激活值（1F1B、全量重计算：第一段存 min(p, m) 个微批次 × 每段的层数 × 每层的输入）；
• 每步的时间 = 计算 / (1 − 气泡) + 张量并行的 all-reduce（NVLink）+ 流水线相邻段的收发 + 数据并行的梯度通信（InfiniBand）。

GPT-3 175B 估计最快的几个组合（全部 27 个组合里 18 个放得下）：

405B 的稠密模型（Llama-3.1 405B 的形状）估计最快的也是 8 × 8 × 16（每块 GPU 32 GB、每步约 24 秒；14 个组合放得下）。规律（也是 Megatron-LM 的经验）：张量并行限制在一台机器的 8 块 GPU 之内（每层都要 all-reduce，只有 NVLink 撑得住）；流水线并行跨机器（只在相邻的段之间传激活，量小），但段越多气泡越大，要用足够多的微批次填满；剩下的 GPU 全给数据并行（每步只通信一次，还能和反向传播重叠）。这个模型没有算注意力的 FLOPs、通信和计算的重叠、延迟项，“秒数”只用来比较组合的相对好坏。

7. 小结 · Summary

1. 训练时每个参数 16 字节，再加上激活值；一块 GPU 放不下大模型，GPU 之间的通信又是主要瓶颈——分布式训练就是在”切开”和”少传”之间找平衡。
2. 数据并行：每块 GPU 一份完整的模型、一部分数据，每步一次梯度 all-reduce；最简单、最有效，但要放得下。
3. 张量并行：MLP 先列并行再行并行、注意力按头切，每层前向 2 次、反向 2 次 all-reduce；只在节点内用。
4. 流水线并行：按层切，气泡 = (p−1)/(m+p−1)；1F1B 省激活显存；异步流水线要暂存权重；用动态规划切层；token 并行让粒度更细。
5. ZeRO：把优化器状态、梯度、参数在数据并行组里切开，显存分别降到 $4\Psi + K\Psi/N_d$、$2\Psi + (2+K)\Psi/N_d$、$(2+2+K)\Psi/N_d$；前两步不增加通信，第三步多 50%。
6. 组合：张量并行在节点内、流水线跨节点、其余给数据并行——“数据并行仍然是王者”。

8. 常见误区 · Common pitfalls

• “数据并行就是把模型复制几份，各训练各的。” 每一步都要把梯度 all-reduce 成同一个平均梯度，各份模型始终一模一样（我们的：和整批差 3e-08）。
• “张量并行切得越多越好。” 每层都要 all-reduce，而且在关键路径上；跨节点时通信和计算一样长（我们的：1.00）。
• “流水线并行就是把层分给不同的 GPU，没有别的代价。” 有气泡；微批次少时大部分 GPU 在等（m = 1 时 75%）。
• “1F1B 比 GPipe 的气泡小。” 一样大（同样的 (p−1)/(m+p−1)）；它省的是激活值的显存。
• “ZeRO 是一种新的优化器。” 它不改变 Adam 的任何一步计算，只改变谁存什么、谁更新什么（我们的：和单卡差 1e-07）。
• “ZeRO 一定会增加通信。” P_os 和 P_os+g 的通信和普通数据并行一样（reduce-scatter + all-gather = all-reduce），只有 P_os+g+p 多 50%。
• “把模型的权重放得下就能训练。” 还要放梯度、优化器状态（共 16 字节 / 参数）和激活值。

9. 练习 · Exercises

练习 1（显存账）。 一个 70B 的稠密模型用混合精度 Adam 训练。(a) 不算激活值，要多少显存？至少几块 80 GB 的 GPU？(b) 用 ZeRO-3 在 64 块 GPU 上，每块要多少？(c) 用 8 路张量并行 + 8 路流水线并行 + ZeRO-1（数据并行 16 路），每块要多少？

(a) 16 × 70e9 = 1120 GB，至少 14 块。(b) 1120 GB / 64 = 17.5 GB。(c) 参数和梯度 4 × 70e9 / 64 = 4.4 GB，优化器状态 12 × 70e9 / (64 × 16) = 0.8 GB，共约 5.2 GB（再加激活值）。

练习 2（列并行的反向）。 用 §3.3 的记号证明：列并行 $W = [W_1, W_2]$ 时，$\partial\ell/\partial X = G_1W_1^\top + G_2W_2^\top$，而 $\partial\ell/\partial W_i = X^\top G_i$ 不需要通信。

$Y = [Y_1, Y_2] = [XW_1, XW_2]$，所以 $\ell$ 通过两条路依赖 X：$\partial\ell/\partial X = \frac{\partial\ell}{\partial Y_1}\frac{\partial Y_1}{\partial X} + \frac{\partial\ell}{\partial Y_2}\frac{\partial Y_2}{\partial X} = G_1W_1^\top + G_2W_2^\top$——两块 GPU 各算一项，要一次 all-reduce。$W_i$ 只出现在 $Y_i = XW_i$ 里，$\partial\ell/\partial W_i = X^\top G_i$，X 和 $G_i$ 第 i 块 GPU 都有，不需要通信。

练习 3（为什么先列后行）。 如果 Megatron 的 MLP 反过来——第一层 A 按行切、第二层 B 按列切——要多少次通信？为什么不行？

第一层行并行：$XA = \sum_i X_iA_i$，要先求和（all-reduce）才能得到完整的 XA，才能做 GeLU——因为 $\mathrm{GeLU}(\sum_i P_i) \ne \sum_i\mathrm{GeLU}(P_i)$。求和之后第二层列并行，输出 $[YB_1, YB_2]$ 是按列切开的，下一层（残差、LayerNorm 需要完整的向量）之前还要一次 all-gather。前向至少两次通信，比”先列后行”的一次多。先列并行，GeLU 的输入每块 GPU 就是完整的那几列，逐元素的 GeLU 可以各算各的。

练习 4（气泡）。 一个模型切成 p = 8 段。(a) 要让气泡低于 10%，至少要几个微批次？(b) 用 GPipe 和 1F1B 时，第一段要同时存几个微批次的激活？(c) 如果全局批大小固定为 512 个序列、数据并行 16 路、每个微批次 1 个序列，m 是多少？气泡呢？

(a) $(p-1)/(m+p-1) < 0.1$ → $7 < 0.1(m + 7)$ → $m > 63$，至少 64 个。(b) GPipe 存 m 份（64）；1F1B 存 p 份（8）。(c) m = 512 / 16 = 32，气泡 7/39 ≈ 18%。要降气泡，就要么减少数据并行（每条流水线分到更多微批次）、要么加大全局批——这就是流水线并行和数据并行之间的取舍。

练习 5（ZeRO 的通信）。 推导：普通数据并行每步每块 GPU 发出 $2\Psi(N-1)/N$ 个参数的数据；ZeRO-2 用 reduce-scatter 和 all-gather 也是这么多；ZeRO-3 为什么是 $3\Psi(N-1)/N$？

环形 all-reduce = 环形 reduce-scatter（每块发出 $\Psi(N-1)/N$）+ 环形 all-gather（再发出 $\Psi(N-1)/N$），共 $2\Psi(N-1)/N$。ZeRO-2：梯度 reduce-scatter（$\Psi(N-1)/N$）+ 更新后参数 all-gather（$\Psi(N-1)/N$），一样多。ZeRO-3：参数平时是切开存的，前向前要 all-gather 一次、反向前还要 all-gather 一次（前向用完就扔了），再加梯度 reduce-scatter，共 3 次，$3\Psi(N-1)/N$；不需要更新后的 all-gather，因为参数本来就切开存。

练习 6（异步流水线的版本差）。 在 1F1B 的稳定阶段，为什么第 s 段（从 1 数起，共 p 段）的一个微批次，前向和反向之间这一段的权重恰好被更新了 p − s 次？

1F1B 里第 s 段在做微批次 j 的前向时，已经有 p − s 个更早的微批次在它后面的各段里”排队”等着回来做反向（这就是 1F1B 让第 s 段同时有 p − s + 1 个微批次在路上的原因）。微批次 j 的反向回到第 s 段之前，这 p − s 个更早的微批次的反向都会先在第 s 段做完，每做完一个就更新一次权重——所以是 p − s 次。最后一段（s = p）前向完马上反向，0 次；第一段 p − 1 次。代码 03 的 3 / 2 / 1 / 0 正是这样。

10. 延伸阅读 · Further reading

• 张量并行：Shoeybi et al. (2019), Megatron-LM；Korthikanti et al. (2022), Reducing Activation Recomputation in Large Transformer Models（序列并行、选择性重计算）。
• 流水线并行：Huang et al. (2019), GPipe；Narayanan et al. (2019), PipeDream；Narayanan et al. (2021), Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM（3D 并行、交错的 1F1B）；Li et al. (2021), TeraPipe。
• ZeRO 与 FSDP：Rajbhandari et al. (2020), ZeRO；Zhao et al. (2023), PyTorch FSDP。
• 专家并行：Lepikhin et al. (2020), GShard；Fedus et al. (2021), Switch Transformers。

11. 下一讲预告 · Next lecture

第 18 讲：FlashAttention 与高效解码。训练时的瓶颈常常在 GPU 之间；而在一块 GPU 之内，注意力的瓶颈其实是读写显存，而不是计算——FlashAttention 把注意力切成小块、在片上内存里算完 softmax，再也不把 n × n 的注意力矩阵写回显存。解码时，一次只生成一个 token，大模型大部分时间在等权重从显存读出来：推测解码（speculative decoding）用一个小模型先猜几个 token、再让大模型一次验证，能少跑几次大模型。