Lecture 9 · 分布式训练：数据、张量、流水线并行与去中心化训练 · Distributed Training

对应 Maps to：Stanford CS324 Winter 2023 Advances in Foundation Models，1/18 客座报告 Ce Zhang（ETH Zürich）：Model Training。Syllabus “Training Methods and Infrastructure” 主题下的两篇：Shoeybi et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism（2019），Hugging Face 博客 The Technology Behind BLOOM Training（2022）。 依据说明：课程网站没有提供 slides 或录像，报告只有一个笼统的题目”Model Training”，具体内容不可考。§7 依据 Ce Zhang 作为作者的三篇公开论文重建：Yuan et al. 2022（Decentralized Training of Foundation Models in Heterogeneous Environments）、Wang et al. 2022（AC-SGD，Fine-tuning Language Models over Slow Networks using Activation Compression with Guarantees）、Tang et al. 2021（1-bit Adam）；这是他 2022 年前后的主要方向，但我不能确定报告讲了其中哪些。§3–§6 依据 Megatron-LM 论文、Rajbhandari et al. 2020（ZeRO）、Narayanan et al. 2021（Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM）、Korthikanti et al. 2022（Reducing Activation Recomputation in Large Transformer Models）、BLOOM 博客、BLOOM 论文和 BigScience 公开的训练说明，以及 PaLM 论文第 4 节。文献里的数字都是我在原文里读到的；§8”2023 年之后”是我补充的。 配套代码 Code：code/lecture09/（numpy + PyTorch；5 个脚本 + 1 个模块。代码 01–04 各 1–4 秒，代码 05 约 55 秒） 预计阅读时间：3.5 小时（含运行代码）

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

0.1 学习目标 Learning goals

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

1. 说清训练一个模型时显存都花在哪里：参数、梯度、优化器状态（混合精度 + Adam 每个参数 16 字节）和激活，并能估算一个模型需要多少张卡。
2. 解释数据并行为什么在数学上等价于用一个更大的 batch，ring all-reduce 怎样工作、为什么每张卡的通信量几乎不随卡数增长，以及数据并行扩展的两个极限。
3. 描述 ZeRO 的三个阶段（切优化器状态、梯度、参数）各省多少显存、多花多少通信，以及它为什么和普通数据并行逐位相同。
4. 推导 Megatron 的张量并行为什么是”第一层按列切、第二层按行切”，每层几次 all-reduce，为什么它几乎总是留在一台机器之内。
5. 计算流水线并行的气泡比例 $(p-1)/(m+p-1)$，并解释 GPipe 和 1F1B 的区别（相同的气泡、不同的激活显存）。
6. 用 BLOOM 的例子解释 3D 并行怎样在显存、气泡、通信和矩阵大小之间折中。
7. 概述 Ce Zhang 等人的去中心化训练：在慢速、异构的网络上怎样分配任务、怎样压缩通信，以及它的代价。
8. 了解 2023 年之后的发展：FSDP 成为标配、Llama 3 的 4D 并行、上下文并行、专家并行、低通信训练。

0.2 先修与本讲范围 Prerequisites & scope

如果你学过本目录下的 CS336，它的第 7、8 讲（并行 1、2）在真实的多 GPU 环境里实现这些并行方式，讲得更细。本讲自成一体：我们只有一个 CPU 进程，所以把每张”卡”表示成 Python 列表里的一个元素，把”发送 / 接收”写成显式的数组拷贝，数清楚每张卡发了多少数据，并逐数字验证每一种并行方式都和单卡训练算的是同一个东西。

0.3 本讲地图 Map

§1  为什么要分布式：显存墙与时间墙；三个切法的总览
§2  显存都花在哪里：模型状态 16Ψ、激活（实测 + 公式）、激活重算               ← 代码 01
§3  数据并行：等价于大 batch；ring all-reduce；扩展的极限；ZeRO                ← 代码 02（以及代码 01 演示 3）
§4  张量并行（Megatron）：按列 / 按行切、按头切、f 与 g、每层 4 次 all-reduce   ← 代码 03
§5  流水线并行：气泡、GPipe、1F1B、阶段之间传什么                             ← 代码 04 演示 1–4
§6  3D 并行：经验法则、BLOOM 的配置、PaLM 的另一种选择                         ← 代码 04 演示 5
§7  去中心化训练（Ce Zhang）：异构网络上的任务分配、通信压缩                    ← 代码 05
§8  2023 年之后：FSDP、4D 并行、上下文并行、专家并行、低通信训练
§9–13  小结、常见误区、练习、延伸阅读、下一讲预告

0.4 记号 Notation（只列本讲新增的，其余见第 1 讲 §0.4）

0.5 关于数字 · About the numbers

讲义里的实验数字都来自 code/lecture09/ 的实际输出（固定种子、PyTorch 单线程，同一台机器上重复运行逐字相同）。本讲的数字有两类：模拟（显存记账、通信量、气泡、异构网络的通信时间）是按公式和简化的代价模型算出来的，用来比较不同做法、理解量级；验证（数据并行、ZeRO、张量并行、流水线并行和单卡训练逐数字对照）是在第 1 讲的 base model 上真的算出来的。有几处结果和我写代码前的预想不一样，讲义按实际结果来写（完整清单在代码 README 的”关于实验结果”里），最值得注意的是：数据并行和单卡的训练轨迹并不逐位相同（§3.1）；简单的局部搜索会卡在平台上（§7.3）；误差反馈在 Adam 上几乎没用（§7.4）。

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

1.1 两堵墙

一个大模型的训练会同时撞上两堵墙。

显存墙。 第 8 讲算过，GPT-3 的 16 位参数是 350 GB，而一张 A100 只有 40 或 80 GB。更糟的是，训练需要的远不止参数（§2）：混合精度 + Adam 每个参数要 16 字节，GPT-3 的”模型状态”是 2.8 TB——至少要 35 张 80 GB 的卡才放得下，还没算激活（代码 01 演示 1）。

时间墙。 即使有一张显存无限大的卡，GPT-3 的 $3.15 \times 10^{23}$ 次浮点运算，按 A100 峰值的 50% 算也要约 $2 \times 10^9$ 秒——64 年。PaLM 540B 在一块 TPU v4 上要六百多年。基础模型能在几周到几个月里训练出来，唯一的原因是成百上千张卡同时在算。

所以分布式训练要回答的问题是：怎样把一次训练切开，分给很多张卡，让它们既放得下、又不互相等待？

1.2 三种切法

一次训练里有三样东西可以切：

                    切什么                     每张卡上有什么                  要通信什么
数据并行（DP）      数据（batch）              一份完整的模型、一部分数据        梯度（all-reduce）
张量并行（TP）      每一层里的矩阵             每一层的一部分                  每层的激活（all-reduce）
流水线并行（PP）    层                         连续的几层                      阶段之间的激活（点对点）

再加上 ZeRO（在数据并行的卡之间切”模型状态”），以及 2023 年以后越来越重要的序列 / 上下文并行（切序列）和专家并行（MoE 的不同专家放在不同的卡上，§8）。真实的大模型训练把几种组合起来——BLOOM 用了数据 × 张量 × 流水线的”3D 并行”，Llama 3 用了”4D”。

这一讲的一个核心论点是：这些切法都不改变算法。 数据并行等价于一个更大的 batch；ZeRO 和普通数据并行逐位相同；张量并行、流水线并行和单卡算的是同一个函数。它们改变的只是”谁来算、谁存什么、谁和谁通信”——所以选择哪种切法是一个纯粹的系统问题，代价用显存、通信量、空闲时间来衡量。代码 02–04 会逐一验证这件事。

2. 显存都花在哪里 · Where the memory goes

2.1 模型状态：每个参数 16 字节

ZeRO 论文（Rajbhandari et al. 2020）把训练时的显存分成两部分：模型状态（model states）和剩余状态（residual states，主要是激活）。混合精度 + Adam 训练时，每个参数要存：

代码 01 演示 1 算了几个模型：GPT-2 1.5B 的 16 位参数只有 3 GB，模型状态却要 24 GB（ZeRO 论文的例子，8 倍）；GPT-3 2.8 TB；PaLM 540B 8.64 TB。BLOOM 的检查点每个参数存 14 字节（bf16 参数 2 + fp32 主参数 4 + Adam 8），176B 参数合计 2.46 TB（2.24 TiB），和训练日志里”完整检查点 2.3TB、只存 bf16 参数 329GB”对得上。

我们的 base model 是纯 fp32 训练，代码 01 演示 2 实测：109,056 个参数，参数 436,224 字节、梯度 436,224 字节、AdamW 的两个矩 872,448 字节——同样是每个参数 16 字节（4 + 4 + 8）。混合精度并不省模型状态，省的是激活和计算时间。

注意优化器状态是大头：12 字节，比参数和梯度加起来（4 字节）还多 3 倍。ZeRO 的第一步就是切它（§3.4）。

2.2 激活：和 batch、序列长度一起增长

前向传播时，每一层都要把一些中间结果存下来，反向传播才能算梯度（第 3 讲 §3 的每一个运算都有自己要存的东西）。这部分叫激活，它和参数量无关，而和 batch 大小、序列长度成正比。

我们可以直接数出来。代码 01 演示 4 用 PyTorch 的 saved_tensors_hooks 拦截 autograd 为反向传播存下的每一个张量（排除模型参数本身），在 7 种配置（宽度、头数、batch、序列长度）上测量，然后用三个特征做最小二乘拟合：

7 个配置的最大相对误差只有 6.5e-04。第一项：每层、每个 token、每个隐藏维度存 17 个数（LayerNorm 的输入、q/k/v、MLP 的中间结果……）；第二项：每层、每个头、每一对位置存 2 个数（注意力的 $s \times s$ 矩阵）；第三项是输出层和交叉熵。

Korthikanti et al.（2022）为 GPT 式的层推导出了同样结构的公式（16 位，单位字节）：

34 字节在 16 位下正好是 17 个数——和我们实测拟合的系数一模一样（$5as^2b$ 那一项包括 softmax 的输出和 dropout 的掩码，我们的模型没有 dropout，所以是 2 个 fp32 数）。

代码 01 演示 5 把这个公式用到大模型上。GPT-3（$h = 12288$，$a = 96$，一个 2048 token 的序列）：

一个序列的激活就和 GPT-3 的 16 位参数差不多大；其中 70% 是注意力的 $s \times s$ 矩阵——这正是第 8 讲 FlashAttention 不存的那部分。

一个值得注意的细节：“激活有多大”非常依赖记账方式。 ZeRO 论文说 GPT-2 1.5B（序列 1K、batch 32）的激活”约 60 GB”——这是按每层每 token 每维 12 个 16 位数的粗略估计（我们重算得 60.40 GB）。用 Korthikanti 的公式，含注意力的 $s^2$ 项是 286.89 GB，去掉它（相当于用 FlashAttention）是 85.56 GB。三种记账相差近 5 倍。结论（激活和模型状态一样是显存大户）不变，但引用具体数字时要先看它怎么算的。

2.3 激活重算：用计算换显存

如果激活放不下，一个通用的办法是激活重算（activation recomputation，也叫梯度检查点 gradient checkpointing）：前向时只存每一层（或每几层）的输入，其余的中间结果算完就扔；反向时，先用存下的输入把这一层的前向重新算一遍，再算梯度。

代码 01 演示 4 在 base model 上用 torch.utils.checkpoint 给每个 Block 做了重算：存下的激活少了 9–17 倍（base model 的配置是 12.3 倍），代价是前向 + 反向的 FLOPs 多了 23.9%。理论上限是 1/3：前向 $2N$ + 反向 $4N$ + 重算 $2N$ = $8N$ vs $6N$。这正是第 8 讲 §7 里 HFU 高于 MFU 的来源：重算的 FLOPs 让硬件忙着，但不算”模型必须做的”。Korthikanti et al. 的贡献之一是选择性重算：只重算占显存多、但计算便宜的那部分（主要就是注意力的 $s \times s$ 矩阵），在 530B 模型上把 MFU 从全部重算时的 42.1% 提高到 54.2%。

2.4 工具箱

把上面的账和后面几节对应起来：

3. 数据并行 · Data parallelism

3.1 最简单的并行：等价于一个大 batch

数据并行（data parallelism, DP）：$N_d$ 张卡各有一份完整的模型；每一步把一个 batch 分成 $N_d$ 份，每张卡在自己那份数据上算前向、反向，得到一份梯度；然后把 $N_d$ 份梯度平均，每张卡用同一个平均梯度更新自己的模型——所有卡的模型始终保持相同。

为什么这样做是对的？损失是对 batch 里所有样本的平均，梯度是线性的：

所以数据并行就是用一个更大的 batch，只是把计算分给了不同的卡。代码 02 演示 3 验证了这一点：在第 1 讲的 base model 上，4 张卡各算 8 个序列、all-reduce 后取平均，和单卡在整个 32 个序列上算的梯度，在同一组参数上只差 1e-7 量级（float32 求和顺序不同带来的舍入）。

但同一个演示还揭示了一个我事先没有想到的现象：如果真的让两个模型分别按”单卡”和”数据并行”训练，两条轨迹并不相同。第 0 步更新之后参数就差了 3.7e-4，40 步后差 3.1e-3（损失曲线仍然几乎重合）。原因是 Adam 按每个分量的 $\sqrt{v}$ 归一化：梯度很小的那些分量，1e-7 的舍入误差相对它们自己并不小，归一化之后就变成了 1e-4 量级的更新差别，再被后面的训练继续放大。

这在真实的大模型训练里是一个实际问题：换一种并行方式、换一个卡数，训练就不再逐位可复现。 如果训练中途出了问题（第 10 讲的损失尖峰），你没法确定”重跑一遍会不会再出同样的问题”。PaLM 特意把整个训练做成”从任何检查点重跑都逐位相同”，原因就在这里（第 10 讲）。

3.2 ring all-reduce

“把 $N_d$ 份梯度加起来，再发给每张卡”——这个操作叫 all-reduce。最直观的做法是参数服务器：所有卡把梯度发给一个中心节点，它加起来再发回去。问题是中心节点要收 $p - 1$ 份、发 $p - 1$ 份，它的网络带宽成了瓶颈，通信时间随卡数线性增长。

ring all-reduce（2017 年由百度引入深度学习，是 NCCL 等通信库的标准算法之一）把 $p$ 张卡排成一个环，把每张卡的数组切成 $p$ 块，分两个半程：

1. reduce-scatter（$p - 1$ 步）：每一步，每张卡把一块发给环上的下一张卡，下一张卡把它加到自己的同一块上。$p - 1$ 步之后，每张卡恰好有一块是所有卡之和。
2. all-gather（$p - 1$ 步）：每张卡把自己那块完整的和沿着环传下去，下一张卡直接覆盖。$p - 1$ 步之后，每张卡都有了全部 $p$ 块。

代码 02 演示 1 一步一步地打印了 4 张卡的过程（每张卡一个长度 8 的数组，切成 4 块）：6 步之后每张卡都得到 1 + 2 + 3 + 4 = 10 倍的 [0..7]，每张卡发出 6 个数——正好是 $2(p-1)/p \times 8$。

关键性质：每一步，每条链路上只传 $1/p$ 的数据，所有链路同时在传。 每张卡一共发送 $2(p-1)/p$ 倍的数组大小——$p$ 从 2 到 16，只从 1 倍涨到 1.875 倍，上限是 2 倍（演示 2）。卡数从 8 加到 1000，每张卡的通信量几乎不变。这是数据并行能扩展到上千张卡的基础。

3.3 数据并行的两个极限

但”通信量不随卡数增长”也意味着通信时间不随卡数减少。代码 02 演示 2 用 α–β 模型算了一个例子：一个 1.3B 参数的模型（16 位梯度 2.6 GB），每张卡 100 Gbit/s 的网络，每一步的 batch 固定为 100 万个 token：

batch 固定时，卡越多，每张卡的计算越少，而梯度有多大就得传多少。8 张卡时通信只占 6%；512 张卡时计算只剩 0.1 秒，通信却要 0.4 秒。实际系统会把通信和反向传播重叠起来（一层的梯度算完就开始传，PyTorch 的 DDP 把梯度分成”桶”来做），这能撑得更久，但终究会被通信追上。

第二个极限来自优化本身：数据并行要用更大的 batch。第 4 讲 §8.2 讲过临界 batch 大小：超过它之后，batch 加倍不再让每一步的进展加倍。卡多到一定程度，要么每张卡分到的数据太少（矩阵太小，第 8 讲 §3.3），要么总 batch 太大、浪费算力。

还有第三个、更根本的限制：每张卡都要放下一整份模型状态。对 GPT-3 这是 2.8 TB——纯数据并行根本训练不了它。这就引出了 ZeRO。

3.4 ZeRO：把模型状态切开

ZeRO（Zero Redundancy Optimizer，Rajbhandari et al. 2020）的观察是：数据并行的 $N_d$ 张卡上存着 $N_d$ 份完全相同的模型状态——巨大的冗余。能不能每张卡只存 $1/N_d$？

关键在于把 all-reduce 拆开用。all-reduce = reduce-scatter + all-gather；reduce-scatter 结束时，每张卡恰好拿到梯度之和的一块。ZeRO 在两个半程之间插进一步：每张卡只用这一块梯度，更新自己负责的那一块参数（只需要这一块的 Adam 状态），然后 all-gather 把更新后的参数传给所有卡。三个阶段：

（代码 01 演示 3，复现了 ZeRO 论文图 1。）ZeRO-1、ZeRO-2 有天花板：参数（或参数 + 梯度）仍然每张卡一整份，卡再多也省不掉——ZeRO-1 在 $N_d = 1024$ 时，一张 80 GB 的卡也只能放下约 20B 参数的模型状态。ZeRO-3 的显存随卡数线性下降，1024 张卡能放下 5 万亿参数的模型状态。

代价是通信。 代码 02 演示 4 在 base model 上模拟了 4 张卡、15 步：

两个结论：(1) ZeRO 的三个阶段和普通数据并行逐位相同（差 0.0）——AdamW 是逐元素的，分段更新和整体更新完全一样，而 reduce-scatter + all-gather 本来就是 all-reduce 做的事。(2) ZeRO-1、ZeRO-2 的通信量和普通数据并行一样（$2(p-1)/p \cdot \Psi$）——省显存不花额外的通信；ZeRO-3 因为参数也切开了，前向、反向之前都要先 all-gather 参数，通信量是 1.5 倍（ZeRO 论文 §7 的结论）。

ZeRO 论文报告：ZeRO-100B 在 400 张 GPU 上训练超过 1000 亿参数的模型，吞吐约 15 PFLOPS；不用模型并行就能训练 130 亿参数的模型（当时最大的 T5 是 110 亿）；微软的 Turing-NLG（170 亿参数）就是用它训练的。PyTorch 后来的 FSDP（Fully Sharded Data Parallel）是 ZeRO-3 思想的官方实现；第 10 讲的 OPT-175B 就是用 FSDP + 张量并行训练的。

4. 张量并行：Megatron-LM · Tensor parallelism

4.1 什么时候需要切开一层

ZeRO-3 把参数切开存，但算的时候还是要把一整层的参数 all-gather 到每张卡上。如果一层本身就大到一张卡放不下（或者算一层的激活放不下），就需要把一层本身切开，让每张卡只算这一层的一部分。这就是张量并行（tensor parallelism，也叫层内模型并行）。

难点在于：切开之后，每张卡的计算还得能拼回同一个结果，而且通信要尽量少。Megatron-LM（Shoeybi et al. 2019）给出了一个对 Transformer 特别简洁的切法。

4.2 MLP：第一层按列，第二层按行

Transformer 的 MLP 是 $Z = \mathrm{GELU}(XA)\,B$，$A$ 把宽度 $h$ 变成 $4h$，$B$ 再变回 $h$。

第一层 $A$ 有两种切法。 按行切（把输入维度切开），每张卡只有 $X$ 的一部分维度，算出 $X_r A_r$；完整的 $XA = \sum_r X_r A_r$。但 GELU 是非线性的：$\mathrm{GELU}(X_1 A_1 + X_2 A_2) \neq \mathrm{GELU}(X_1 A_1) + \mathrm{GELU}(X_2 A_2)$，所以必须在 GELU 之前先 all-reduce 一次。按列切（把输出维度切开）：$A = [A_1, A_2]$，每张卡算 $\mathrm{GELU}(XA_r)$——GELU 是逐元素的，各算各的，不需要通信。Megatron 选了按列切。

第二层 $B$ 按行切：第一层的输出已经按列分好了（第 $r$ 张卡有 $Y_r = \mathrm{GELU}(XA_r)$），正好是第二层输入的第 $r$ 段；每张卡算 $Y_r B_r$，最后 $Z = \sum_r Y_r B_r$——一次 all-reduce。

代码 03 演示 1 在 base model 第一层的 MLP 上（float64）验证：Megatron 的切法和单卡的差在 1e-15 量级（$t = 2, 4, 8$）；而第一层按行切、GELU 之前不同步，误差是 0.71–0.89——比输出本身的典型大小（0.355）还大，结果完全错了。

4.3 注意力：按头切

多头注意力天然可以切：不同的头在注意力内部没有任何交互。把 $q, k, v$ 的投影按列切（第 $r$ 张卡只算自己那几个头的 $q, k, v$），每张卡做完整的注意力，输出投影按行切，最后一次 all-reduce。代码 03 演示 2：base model 的 4 个头分到 1、2、4 张卡上，和单卡的差都在 1e-16 量级。这要求头数能被 $t$ 整除——大模型的头数常常是 8 的倍数（GPT-3 96 个，PaLM 540B 48 个）。

4.4 f 和 g：每层 4 次 all-reduce

反向传播怎么办？代码 03 演示 4 说明：每张卡上那一片权重的梯度，正好是单卡梯度的对应切片（差 1e-18 量级）——不需要额外通信。需要通信的是输入 $X$ 的梯度：每张卡都用了完整的 $X$，各自只算出 $X$ 的梯度的一部分（来自自己那几列），要加起来。

Megatron 把这写成一对”共轭”的算子：

• $f$：前向恒等（把 $X$ 复制给每张卡），反向 all-reduce（把 $X$ 的梯度加起来）；
• $g$：前向 all-reduce（把各卡的输出加起来），反向恒等。

MLP 和注意力各有一对 $f$、$g$，所以每个 Transformer 层前向 2 次、反向 2 次，一共 4 次 all-reduce。代码 03 演示 3 把整个 base model 切成 $t = 1, 2, 4$ 份：验证集 bit/字符都是 0.408576（和单卡差 1e-8 以下），每个序列的前向 4 次 all-reduce（2 层 × 2）。

4.5 为什么张量并行留在一台机器里

张量并行每次 all-reduce 的是激活，大小 = token 数 × 宽度，不随卡数减少，而且每一层都要做。代码 03 演示 5 算了 GPT-3 的一层（一个 2048 token 的序列，$t = 8$）：一次 all-reduce 50.3 MB，每张卡每层（前向 + 反向 4 次）发 352 MB；这一层的计算切成 8 份、每张卡按 A100 峰值的 50% 算是 17.84 ms。

（两种带宽都来自 Megatron 论文的硬件：DGX-2H 服务器内 GPU 之间 300 GB/s，服务器之间 8 块 InfiniBand 网卡共 100 GB/s。）在一台机器内部，通信只是计算的 7%；跨机器，通信是计算的 3 倍多——卡大部分时间在等数据。所以实践中张量并行几乎总是放在一台机器的 8 张卡之内。

Megatron-LM 论文的结果：在 512 张 V100（32 台 DGX-2H）上训练 83 亿参数的 GPT-2，整个应用持续 15.1 PFLOPS；以单卡训练 12 亿参数模型的 39 TFLOPS（峰值的 30%）为基准，扩展效率 76%（8 路模型并行时 77%）。这个 8.3B 模型在 WikiText103 上困惑度 10.81（此前最好 15.79），LAMBADA 准确率 66.51%（此前最好 63.24%）。论文还有一个对架构的发现：BERT 类模型超过 3.36 亿参数后会退化，把 LayerNorm 和残差连接的位置调整一下（即后来的 pre-LN 写法，第 3 讲 §3.6）就能继续扩展。

序列并行（Korthikanti et al. 2022）：张量并行切开了注意力和 MLP 的大部分激活，但 LayerNorm、dropout 这些部分在每张卡上仍然是完整的一份（$10sbh$ 那一项，§2.2 的表）。把它们沿序列方向切开，整层的激活就都除以 $t$ 了（§2.2 表里的第三行）。

5. 流水线并行 · Pipeline parallelism

5.1 按层切，以及气泡

张量并行不能跨机器，那跨机器怎么切模型？流水线并行（pipeline parallelism）：把模型的层分成 $p$ 段（“阶段”，stage），第 1 张卡算前几层，把激活传给第 2 张卡……反向时梯度沿相反的方向传回来。阶段之间只传切口处的激活——点对点通信，比 all-reduce 便宜得多，适合放在机器之间。

问题是等待。代码 04 演示 1 画了 4 个阶段的时间线（每格一个时间单位，F = 前向，b = 反向，·· = 空闲）。不切 micro-batch 时：

阶段 0 │F0 F0 F0 F0 ·· ·· ·· ·· ·· ·· ·· ·· ·· ·· ·· ·· … b0 b0 b0 b0 b0 b0 b0 b0
阶段 1 │·· ·· ·· ·· F0 F0 F0 F0 ·· ·· ·· ·· ·· ·· ·· ·· … ·· ·· ·· ·· ·· ·· ·· ··
阶段 2 │·· ·· ·· ·· ·· ·· ·· ·· F0 F0 F0 F0 ·· ·· ·· ·· …
阶段 3 │·· ·· ·· ·· ·· ·· ·· ·· ·· ·· ·· ·· F0 F0 F0 F0 b0 b0 b0 b0 b0 b0 b0 b0 …

任何时刻只有一个阶段在干活：75% 的时间是空闲的——4 张卡只干了 1 张卡的活。这些空闲叫气泡（bubble）。

5.2 GPipe：切成 micro-batch

GPipe（Huang et al. 2019）的办法：把一个 batch 切成 $m$ 个 micro-batch，阶段 0 算完第一个 micro-batch 交给阶段 1 之后，马上开始算第二个。$m = 4$ 时：

阶段 0 │F0 F1 F2 F3 ·· ·· ·· ·· ·· ·· ·· ·· ·· b0 b0 b1 b1 b2 b2 b3 b3
阶段 1 │·· F0 F1 F2 F3 ·· ·· ·· ·· ·· ·· b0 b0 b1 b1 b2 b2 b3 b3 ·· ··
阶段 2 │·· ·· F0 F1 F2 F3 ·· ·· ·· b0 b0 b1 b1 b2 b2 b3 b3 ·· ·· ·· ··
阶段 3 │·· ·· ·· F0 F1 F2 F3 b0 b0 b1 b1 b2 b2 b3 b3 ·· ·· ·· ·· ·· ··

气泡降到 42.9%；$m = 8$ 时 27.3%。气泡只出现在开头（后面的阶段等第一个 micro-batch 到达）和中间（前向做完，等第一个反向传回来）。代码 04 演示 2 的模拟结果和公式逐个相等：

要让气泡小，$m$ 要远大于 $p$：$p = 8$ 时 $m = 64$ 才把气泡压到 10% 以下。（Narayanan et al. 2021 把它写成”气泡时间 ÷ 理想时间 = $(p-1)/m$“，是同一件事。）但 $m$ 不能无限大：batch 总大小有限，切得越碎，每个 micro-batch 越小，矩阵越”瘦”，硬件利用率越低（第 8 讲 §3.3）。

5.3 1F1B：同样的气泡，更少的激活

GPipe 先做完所有 $m$ 个前向，再做反向——每个阶段要同时存 $m$ 个 micro-batch 的激活。1F1B（PipeDream-Flush，Narayanan et al. 2021）换了一个顺序：第 $s$ 个阶段先做 $p - s - 1$ 个前向”热身”，然后一个前向、一个反向交替。每做完一个反向，那个 micro-batch 的激活就可以扔掉。

代码 04 演示 3（$p = 4$，$m = 16$）：GPipe 每个阶段同时存 16 份激活；1F1B 第 0 个阶段最多存 4 份，最后一个阶段 1 份——和 $m$ 无关。两者的总时间完全相同（57 个时间单位）：1F1B 不减少气泡，减少的是激活显存。这就是 Megatron、DeepSpeed 默认用 1F1B 的原因。

Narayanan et al. 还提出了交错（interleaved）调度：每张卡不是负责连续的一段层，而是负责好几段不连续的层（$v$ 个”虚拟阶段”），气泡缩小到 $1/v$，代价是通信多了约 $v$ 倍。

5.4 流水线 = 梯度累加

流水线并行也不改变算法。代码 04 演示 4 把 base model 切成两个阶段（嵌入 + 第 1 个 Block；第 2 个 Block + 输出层），batch 32 切成 4 个 micro-batch，每个 micro-batch 的损失除以 4 再反向、梯度累加：和不切时整个 batch 的梯度差 1.8e-08。

阶段之间每一步传的数据是 4 × 8 × 96 × 64 个 fp32 数 = 786,432 字节（前向），反向同样大小。在这个玩具模型上它比参数（436,224 字节）还大；但对大模型，一个切口的激活（token 数 × 宽度）远小于整个模型的参数——流水线并行的通信量小、而且是点对点的，所以适合放在带宽较低的机器之间。

6. 3D 并行：把它们组合起来 · Putting it together

6.1 经验法则

Narayanan et al. 2021 系统地研究了三种并行怎样组合，结论可以概括成三条：

1. 张量并行在一台机器之内，份数不超过一台机器的卡数（DGX A100 是 8）——它的通信最多，需要 NVLink。
2. 流水线并行跨机器——点对点通信，对带宽要求低；用它把模型切到放得下为止。
3. 数据并行在最外层——模型放下之后，用它扩展到更多的卡（配合 ZeRO-1 切优化器状态）。

按这个方法，他们在 3072 张 A100 上训练了一个 1 万亿参数的模型，吞吐 502 PFLOP/s，每张卡 163 TFLOP/s——峰值的 52%。

6.2 BLOOM 的配置

BLOOM 176B（BigScience，2022）是一个公开了全部工程细节的例子。根据 Hugging Face 的博客、BLOOM 论文和 BigScience 的训练说明：

• 硬件：法国 Jean Zay 超算上 48 个节点 × 8 张 A100 80GB = 384 张卡（另有 4 个备用节点）；节点内 NVLink，节点之间 Omni-Path 网络。
• 并行：Megatron-DeepSpeed，张量并行 $t = 4$ × 流水线并行 $p = 12$ × 数据并行 $d = 8$ = 384；数据并行上用 ZeRO-1。
• 其他：bf16 混合精度（第 10 讲会讲为什么不用 fp16）；Megatron 的融合 kernel（LayerNorm、“缩放 + 掩码 + softmax”、偏置 + GELU——正是第 8 讲 §4.3 说的那类融合）。
• 结果：最快的配置每张卡 156 TFLOP/s（峰值的一半）；约 3.5 个月，1,082,990 个计算小时，366B 个 token。

为什么是 4 × 12 × 8？代码 04 演示 5 把 384 张卡的几种组合列出来（全局 batch 2048 个序列、micro-batch 2，数据并行上用 ZeRO-1）：

• 不用张量并行（$t = 1$）时，要么放不下，要么流水线很长、每条流水线分到的 micro-batch 少，气泡在 24%–26%。
• BLOOM 的 4 × 12 × 8：每张卡的模型状态约 20 GB，给激活留出大半张卡；气泡约 8%；张量并行 4 份留在一台机器（8 张卡）里。
• 8 × 6 × 8 的气泡更小（3.8%），但张量并行用满了一台机器，每层的通信更多、每张卡的矩阵更小。

没有”最优”的公式答案——要在显存、气泡、通信、矩阵大小之间折中，最后靠实测吞吐来选。BigScience 的训练说明里写着，他们试了几种 micro-batch 大小，选了实测最快的那一个。

6.3 另一种选择：PaLM 不用流水线

PaLM 540B（第 10 讲的主角）走了另一条路：完全不用流水线并行。在 Google 的 TPU v4 Pod 里，3072 块芯片之间有高带宽的专用互联，所以 PaLM 在每个 Pod 内部用 12 路模型并行（张量并行）× 256 路全分片数据并行（相当于 ZeRO-3 / FSDP），两个 Pod 之间再做 2 路数据并行——由 Pathways 系统协调，跨 Pod 时每对主机每步交换约 1.3 GB 的梯度，突发带宽合计 81 Tbps，达到完美弱扩展的 97%。没有流水线，就没有气泡；代价是要求极好的互联。这说明 3D 并行的”经验法则”本身取决于硬件：网络越好，越可以少用流水线。

6.4 混合专家与专家并行

第 3 讲 §8.1 提到，混合专家（MoE）模型每个 token 只激活一部分参数，计算量省了，但所有专家的参数都要存在某处。专家并行把不同的专家放在不同的卡上，每个 token 被路由到它的专家所在的卡——这需要一种新的通信模式（all-to-all：每张卡给每张卡发一些 token）。MoE 在 2023 年之后重新成为主流（§8），专家并行也就成了第四个维度。

7. 去中心化训练：Ce Zhang 的问题 · Decentralized training

7.1 动机：算力在哪里

前面所有的方法都假设卡在同一个数据中心里，用每秒几百 Gbit 到几 Tbit 的网络连着。但世界上大量的 GPU 并不在这样的地方：分散在各个大学、各个云区域、闲置的工作站里，彼此之间是普通的互联网——可能只有 1 Gbit/s，延迟几十到几百毫秒。如果能在这样的网络上训练基础模型，训练它的门槛就会低得多。

这是 Ce Zhang（当时在 ETH Zürich）和合作者在 2021–2022 年的一个主要方向。困难很明显：§3–§5 的每一种并行都要大量通信，慢速网络上通信会吞掉一切。

7.2 任务分配：把重的通信放在快的链路上

Yuan et al.（2022，Decentralized Training of Foundation Models in Heterogeneous Environments，作者包括 Tri Dao、Percy Liang、Chris Ré 和 Ce Zhang）的思路是：网络是异构的——有的卡之间很快（同一个地区、同一个机房），有的很慢（跨大洲）。数据并行和流水线并行的通信量、通信模式不同，把通信量大的任务分配到快的链路上，就能大幅减少通信时间。

他们把训练切成”小任务”（tasklet：一个流水线阶段 × 一个 micro-batch 的前向和反向），把”哪张卡做哪个任务”写成一个优化问题，用两层的代价模型来评估一个分配：上层是把卡分成数据并行组（一个平衡图划分问题），下层是把这些组串成流水线（一个开环旅行商问题）。他们发现标准的图划分启发式（Kernighan–Lin）在这个目标上不起作用，于是设计了专门的局部搜索，配合遗传算法。

实验用 64 张 V100、GPT-3 1.3B 规模的模型，五种场景：数据中心按需实例、数据中心竞价实例、两个数据中心、美国四个地区、全球八个地区（跨地区 0.3–1.3 Gbit/s、延迟 10–250 ms）。最极端的全球场景下，他们的调度比 Megatron 快 4.8 倍、比 DeepSpeed 快 3.6 倍；和在数据中心里训练相比，去中心化只慢 1.7–3.5 倍。

7.3 迷你实验：32 张卡、4 个地区

代码 05 演示 1 做了一个简化版：32 张卡分在 4 个地区（地区内 10 Gbit/s；地区之间 0.3–1.3 Gbit/s、延迟 10–250 ms），训练一个 1.3B 的模型，切成 4 个流水线阶段 × 8 个数据并行副本，全局 batch 1024 个序列。代价模型（dist_sim.iteration_cost）：数据并行部分 = 每个阶段的 8 张卡 all-reduce 0.65 GB 梯度、受环上最慢的链路限制，取最慢的阶段；流水线部分 = 每个副本沿 4 个阶段传激活（每个切口 1.07 GB，前向反向各一次），取最慢的副本。

几个观察：

1. 分配的好坏差 4 倍：随机分配 77.9 秒，最好的分配 19.0 秒。
2. 哪种通信更重，取决于具体的数字。 我原本以为数据并行的梯度同步最重、应该放在地区内；但在这个设置里，流水线要传的激活更多（每个切口 1 GB、前向反向各一次、3 个切口串行），所以把流水线留在地区内、让数据并行跨地区更好（19.1 vs 46.8 秒）。换一个 batch 大小或模型大小，结论可能反过来——这正是需要一个代价模型来搜索、而不是凭直觉的原因。
3. 简单的局部搜索会卡住。 “随机交换两张卡，代价变小就接受”在 39 秒左右停住，迭代再多也出不来（我在练习用的副本里试过 10 万次交换，结果一样）。原因是代价取的是最慢的那个阶段 / 副本：大多数交换根本不碰最慢的那一个，代价不变——搜索停在一大片平台上。允许”代价不变也接受”之后，最好的一次重启找到了 19.0 秒，略好于最好的手工方案（找到的分配也是”24 条流水线边全在地区内”）；但 5 次重启里有 3 次仍然很差。这和论文的经验一致：这个优化问题需要专门设计的搜索。
4. 加上计算，差距缩小。 只看通信，最好的分配也比数据中心慢 7.8 倍；但每次迭代的计算（32 张 V100 按峰值的 40% 算）要 10.2 秒，加上之后只慢 2.3 倍（随机分配慢 7.0 倍）——正好落在论文报告的 1.7–3.5 倍之间。

7.4 压缩通信

另一个手段是少传。代码 05 演示 2 在 base model 上做 4 张卡的数据并行，每张卡发出的梯度先压缩，从随机初始化训练 300 步：

(a) AdamW + 梯度裁剪（第 1 讲的训练配方）

(b) 带动量的 SGD，不裁剪

这里的误差反馈（error feedback）是压缩通信的标准技巧：每张卡记住上一次”没发出去的部分”（压缩前 − 压缩后），下一次加回到新梯度上再压缩——被压掉的信息没有丢，只是晚一点发出去。

结果和我写代码前的预想有好几处不同：

• fp16 免费（1.5070 vs 1.5065），通信量减半。int8 有可见的损失：整个向量只用一个缩放因子时，缩放因子由最大的那个梯度决定，大量很小的梯度被舍入成 0（实际系统会按块分别量化）。
• top-1% 很伤：300 步里每个分量平均只被发出 3 次。误差反馈在带动量的 SGD 上有明显帮助（3.11 → 2.79），在 AdamW 上几乎没用（2.82 → 2.77）。
• 只发符号时，AdamW 不用误差反馈反而好得多（1.74 vs 2.69）。Adam 本来就按每个分量的大小归一化，更新的方向主要由符号决定，只发符号丢的信息不多；而误差反馈累积的残差会被 Adam 的归一化和梯度裁剪扭曲。

这个”误差反馈在 Adam 上不灵”的现象正是 1-bit Adam（Tang et al. 2021，Ce Zhang 是作者之一）要解决的问题：已有的误差补偿压缩只对 SGD 一类的优化器有效，对 Adam 这种梯度的非线性函数不行。他们的观察是 Adam 的方差项在热身之后会变得稳定，于是先用普通 Adam 热身，然后把方差冻结成一个固定的预条件器，再做 1 bit 压缩——通信量少到 1/5，BERT-Large 预训练的吞吐提高到 3.3 倍，收敛速度和不压缩的 Adam 相同。

流水线并行要传的是激活，压缩它更难：激活的误差会直接改变前向的结果。AC-SGD（Wang et al. 2022）的想法是压缩激活的变化量而不是激活本身——同一个样本在相邻两次见到时，激活变化不大，变化量可以压得很狠。他们把激活压到 2–4 bit，在慢速网络上端到端快了最多 4.3 倍，模型质量不变；再加上梯度压缩，最多快 4.9 倍。

一个实际的例子：2022 年 11 月，Together 发布的 GPT-JT（在 GPT-J-6B 的基础上继续训练 35.3 亿个 token）就是在 1 Gbit/s 的普通网络上，用去中心化的训练方法训练的。

7.5 去中心化的其他难题

本讲只讨论了”慢”和”异构”。真实的去中心化训练还要面对：卡随时可能离开或加入（竞价实例被收回、志愿者关机）；不同的卡算得快慢不一（流水线里最慢的阶段决定整体速度）；如果参与者互不信任，还有恶意节点发送错误梯度的问题（第 17 讲讲投毒）。这些都是开放的研究问题。

8. 2023 年之后 · Since 2023

CS324 Winter 2023 开课时，Megatron + DeepSpeed 的 3D 并行是训练大模型的主流配方。此后的变化可以概括为：切法越来越多，工具越来越成熟，规模又上了一个数量级。这里只列和本讲直接相关、比较确定的几条。

• FSDP / ZeRO-3 成为标配。 PyTorch 的 FSDP（以及后来的版本）让”把参数、梯度、优化器状态全部切开”成为几行代码的事，中等规模的模型（几十亿到几百亿参数）常常只用 FSDP 就够了，不需要张量或流水线并行。
• 4D 并行：Llama 3 405B（2024）。 在最多 16K 张 H100 上训练，用张量并行（8）、上下文并行（CP）、流水线并行（16）和数据并行（FSDP）四个维度。论文的表 4：8,192 张卡时数据并行 64 份，16,384 张卡时 128 份；训练长上下文时上下文并行 16 份、数据并行减到 8 份；每次迭代 1600 万个 token，每张卡 380–430 TFLOP/s，BF16 MFU 38%–43%。
• 上下文 / 序列并行。 上下文长到 128K 甚至更长时（第 8 讲 §8.3），一个序列的激活一张卡放不下，就把序列切到多张卡上；注意力需要所有位置的 $K, V$，于是有了 ring attention 一类的方法——让 $K, V$ 的块沿着卡组成的环传递，每张卡用在线 softmax（第 8 讲 §5.2）把各块的结果合并。
• 专家并行与 MoE 的回归。 2023 年以后不少最强的开放模型是 MoE（第 3 讲 §9），专家并行（all-to-all 通信）成为第五个维度。例如 DeepSeek-V3（2024 年底）在 2048 张 H800 上训练，专门设计了让通信和计算重叠的流水线调度，并用 FP8 训练（第 10 讲）。
• 低通信训练。 去中心化的方向继续发展：DiLoCo（Douillard et al. 2023）一类方法让每个”岛”本地训练几百步、再同步一次，通信量减少几百倍；2024 年起出现了由分布在多个大洲的算力共同训练的 100 亿参数级别的开放模型。§7 的问题——慢速网络上能不能训练基础模型——的答案越来越接近”能”。
• 编译器自动切分。 JAX / XLA 的 GSPMD（PaLM 用的就是它）以及 PyTorch 的 DTensor 让用户只写”这个张量按哪个维度切”，由编译器插入通信。“手写每一种并行”逐渐变成”声明切分方式”。

哪些结论仍然成立？ 显存账（16Ψ + 激活）、ring all-reduce 的通信量、Megatron 的切法、气泡公式、“张量并行留在机器内”——这些是今天所有大规模训练的基础。被修正或补充的：“3D”不够了，上下文和专家成了新的维度；“流水线必不可少”取决于互联（PaLM 不用，很多 FSDP 训练也不用）；规模从几百张卡涨到几万张卡，故障成了和通信同样重要的问题——Llama 3 的 54 天里有 466 次中断，这是第 10 讲的主题。

9. 小结 · Summary

1. 两堵墙：显存（GPT-3 的模型状态 2.8 TB，至少 35 张 80 GB 的卡）和时间（单卡 64 年）。分布式训练要切开计算、切开存储，又不让卡互相等待。
2. 显存账：混合精度 + Adam 每个参数 16 字节（2 + 2 + 12，优化器状态是大头）；激活 $\approx sbh(34 + 5as/h)$ 字节/层（我们实测拟合出完全相同的系数），其中注意力的 $s^2$ 项在长序列时占大头；激活重算用约 1/3 的额外计算换显存。
3. 所有切法都不改变算法：数据并行 = 大 batch；ZeRO 和普通数据并行逐位相同；张量并行、流水线并行和单卡算同一个函数。但浮点舍入会让不同并行方式的训练轨迹慢慢分开。
4. 数据并行 + ring all-reduce：每张卡发 $2(p-1)/p$ 倍的梯度，几乎不随卡数增长——但也不随卡数减少，卡多了通信占比上升；batch 也不能无限大。
5. ZeRO：切优化器状态（ZeRO-1）、梯度（ZeRO-2）、参数（ZeRO-3）；ZeRO-1/2 不增加通信，ZeRO-3 多 50%。$\Psi = 7.5\text{B}$、64 张卡：120 → 31.4 → 16.6 → 1.88 GB。
6. 张量并行（Megatron）：MLP 第一层按列、第二层按行；注意力按头；每层 4 次 all-reduce（$f$、$g$）；通信量 ∝ 激活，只适合机器内的高速互联。
7. 流水线并行：气泡 $(p-1)/(m+p-1)$，$m \gg p$ 才能小；1F1B 气泡相同、激活最多存 $p$ 份；阶段之间点对点通信、量小，适合跨机器。
8. 3D 并行：张量在机器内、流水线跨机器、数据在最外层。BLOOM：4 × 12 × 8 = 384 张 A100，每张卡约 20 GB 模型状态，气泡约 8%，156 TFLOP/s。PaLM 靠高速互联完全不用流水线。
9. 去中心化：在异构的慢速网络上，任务分配差 4 倍（把重的通信放在快的链路上），通信压缩再省几倍；加上计算后只比数据中心慢 2–3 倍。压缩和优化器相互影响：误差反馈在 Adam 上不灵，1-bit Adam 先冻结方差再压缩。

10. 常见误区 · Common pitfalls

1. “175B 参数 × 2 字节 = 350 GB，所以 5 张 80 GB 的卡就能训练 GPT-3。” 训练要存参数、梯度、优化器状态（16 字节/参数，2.8 TB）和激活。
2. “数据并行让训练快 N 倍。” 它让每一步快（最多 N 倍），但每一步的 batch 也大了 N 倍。超过临界 batch 后，每一步的进展不再随 batch 线性增长；卡多了通信占比也会上升。
3. “ring all-reduce 的时间随卡数减少。” 每张卡要发的数据几乎不变（上限 2 倍梯度大小），时间也几乎不变——随卡数减少的是计算。
4. “ZeRO 是一种近似 / 会改变训练结果。” ZeRO 和普通数据并行逐位相同；它只改变谁存什么。ZeRO-1、ZeRO-2 甚至不增加通信。
5. “张量并行可以随便跨机器。” 它每层都要 all-reduce 激活，跨机器时通信会是计算的几倍。
6. “流水线并行把 GPU 用满了。” 不切 micro-batch 时，4 个阶段的流水线 75% 的时间在空闲；要 $m \gg p$ 才能把气泡压下去。1F1B 也不减少气泡，只减少激活显存。
7. “分布式训练和单卡训练逐位相同。” 数学上等价，但浮点求和顺序不同，Adam 会把舍入误差放大，训练轨迹会分开。想要逐位可复现，需要专门设计（PaLM，第 10 讲）。
8. “梯度压缩 + 误差反馈总能接近不压缩。” 误差反馈的理论保证针对 SGD 一类的优化器；在 Adam 上它可能几乎没用，甚至有害（§7.4）。压缩要和优化器一起设计。
9. “去中心化训练只是更慢的数据中心训练。” 它需要不同的设计：任务分配、通信压缩、容错。好的分配和差的分配可以差好几倍。

11. 练习 · Exercises

练习 1（显存账） 一个 7B 参数的模型，混合精度 + Adam，用 8 张 80 GB 的卡做数据并行。分别计算普通数据并行、ZeRO-1、ZeRO-2、ZeRO-3 每张卡的模型状态。哪些放得下（先不算激活）？

总模型状态 $16 \times 7\text{B} = 112$ GB。用 zero_bytes_per_device（$K = 12$，$N_d = 8$）：普通数据并行 112 GB（放不下）；ZeRO-1：$4\Psi + 12\Psi/8 = 28 + 10.5 = 38.5$ GB；ZeRO-2：$2\Psi + 14\Psi/8 = 14 + 12.25 = 26.25$ GB；ZeRO-3：$16\Psi/8 = 14$ GB。ZeRO-1/2/3 都放得下模型状态，剩下的显存留给激活。这也是为什么 7B 级别的模型常常只用 FSDP / ZeRO 就够了，不需要张量并行。

练习 2（all-reduce 时间） 1.3B 参数的 16 位梯度（2.6 GB），8 张卡，用 α–β 模型（延迟 5 µs）分别算：(a) 每张卡 25 GB/s 的机器内互联；(b) 每张卡 10 Gbit/s（1.25 GB/s）的普通网络。数据并行的每一步要多花多少时间？

环形 all-reduce：$2(p-1)\alpha + 2(p-1)/p \cdot n/\beta$。(a) $\approx 0.182$ 秒；(b) $\approx 3.64$ 秒（延迟项只有 70 µs，可以忽略）。慢 20 倍的网络让每一步多出几秒的通信——如果每一步的计算只有一两秒，训练就被通信主导了。这就是 §7 的去中心化训练要解决的问题。

练习 3（气泡） BLOOM 的流水线 $p = 12$。(a) 每条流水线 128 个 micro-batch 时气泡多大？(b) 要把气泡压到 5% 以下，至少要多少个 micro-batch？在全局 batch 2048、数据并行 8 份的条件下，这意味着 micro-batch 大小是多少？

(a) $11 / (128 + 11) = 7.9\%$。(b) $(p-1)/(m + p - 1) \le 0.05$ ⟺ $m \ge 19(p-1) = 209$。每条流水线分到 $2048 / 8 = 256$ 个序列，$m = 209$ 意味着 micro-batch 不到 2 个序列——只能取 1，$m = 256$，气泡 4.1%。代价是每个 micro-batch 只有一个序列，矩阵更”瘦”，每个阶段的计算效率下降。这正是 BigScience 实测比较 micro-batch 1 和 2 的原因。

练习 4（为什么按列切） 用自己的话解释：为什么 Megatron 把 MLP 的第一层按列切、而不是按行切？如果激活函数是线性的（比如没有激活函数），按行切还需要在中间同步吗？

按行切时每张卡只得到 $XA$ 的一个”部分和” $X_r A_r$，要先加起来才能过 GELU，因为 GELU 是非线性的。按列切时每张卡得到 $XA$ 的完整的几列，GELU 逐元素作用，可以直接算。如果中间没有非线性，$Z = XAB = \sum_r X_r A_r B$，按行切第一层之后可以直接乘 $B$ 再最后加起来，中间不需要同步——非线性是”必须先同步”的根本原因。

练习 5（1F1B 的显存） 用 simulate_pipeline 的思路解释：为什么 1F1B 的第 $s$ 个阶段同时最多存 $p - s$ 份激活？为什么 GPipe 每个阶段都要存 $m$ 份？

1F1B 的第 $s$ 个阶段先做 $p - s - 1$ 个前向热身，然后每做一个前向就做一个反向。热身结束时它存着 $p - s - 1$ 份，再做一个前向变成 $p - s$ 份，接着做一个反向扔掉一份……稳定在 $p - s$。热身的个数正好是”第一个反向从最后一个阶段传回来之前，它能先做的前向个数”。GPipe 先做完所有 $m$ 个前向才开始反向，所以每个阶段在中间时刻存着全部 $m$ 份。

练习 6（张量并行的通信量） GPT-3 的一层，一个 micro-batch 有 8 个 2048 token 的序列，张量并行 $t = 8$，16 位。一次 all-reduce 的数据有多大？每层前向 + 反向一共几次？如果改成 micro-batch 1 个序列，通信量和计算量各怎样变？

一次 all-reduce：$8 \times 2048 \times 12288 \times 2$ 字节 ≈ 0.40 GB；每层 4 次（前向 2、反向 2）。micro-batch 变成 1 个序列，通信量和计算量都变成 1/8——两者都正比于 token 数，比值不变。所以张量并行的”通信 ÷ 计算”主要由宽度和链路带宽决定（计算 ∝ token × $h^2$，通信 ∝ token × $h$，比值 ∝ $1/h$）：模型越宽，张量并行越划算。

练习 7（改一行代码，预测结果） 把代码 02 演示 4 的卡数 $p$ 从 4 改成 8。预测：(a) 参数与普通 DP 的差；(b) 每张卡常驻的数（以 Ψ 为单位）；(c) 每张卡每步发出 ÷ Ψ。

(a) 仍然是 0.0——卡数不影响”分段更新 = 整体更新”。(b) 普通 DP $4\Psi$；ZeRO-1 $2\Psi + 2\Psi/8 = 2.25\Psi$；ZeRO-2 $\Psi + 3\Psi/8 = 1.375\Psi$；ZeRO-3 $4\Psi/8 = 0.5\Psi$。(c) 普通 DP 和 ZeRO-1/2：$2 \times 7/8 = 1.75$；ZeRO-3：$3 \times 7/8 = 2.625$。注意：演示 4 的 batch 是 32、每张卡 8 个序列，改成 8 张卡时要把每张卡的切片改成 4 个序列，否则下标会越界。

练习 8（气泡和前向 / 反向的比例） 在代码 04 里用 simulate_pipeline(4, 4, "gpipe", t_f, t_b) 试 $(t_f, t_b) = (1, 1), (1, 2), (1, 3)$，气泡比例变吗？为什么？

三种情况气泡都是 3/7 = 42.86%。GPipe 的时间线是”前向部分”和”反向部分”拼起来的，每部分都是 $(m + p - 1)$ 个时间片、其中每个阶段忙 $m$ 个，所以气泡比例 $(p-1)/(m+p-1)$ 和 $t_f$、$t_b$ 的比例无关。（如果阶段之间的计算量不均衡——比如第一个阶段多了嵌入层、最后一个阶段多了输出层——就不再是这个公式了：最慢的阶段决定整条流水线的速度。BLOOM 把嵌入层也按一个 Block 来分配阶段，就是为了均衡。）

练习 9（FlashAttention 与激活） 用 Korthikanti 的公式，GPT-3 的一层、一个 2048 token 的序列：(a) 不重算时每层的激活多大？(b) 如果用 FlashAttention（不存 $s \times s$ 矩阵，近似地去掉 $5as/h$ 项），每层多大？96 层呢？

(a) $sbh(34 + 5as/h) = 2048 \times 12288 \times (34 + 80) ≈ 2.87$ GB。(b) $34 sbh ≈ 0.86$ GB 每层，96 层约 82 GB——比 275 GB 少了 70%。FlashAttention 不仅让注意力更快（第 8 讲），还省掉了激活里最大的一块；这使得很多训练可以只做选择性重算甚至不重算，MFU 因此更高。

练习 10（去中心化的数据并行） 一个阶段的梯度 0.65 GB，8 张卡分在不同的地区，环上最慢的链路 1 Gbit/s、延迟 100 ms。一次 ring all-reduce 要多久？如果用只发符号的压缩（每个数 1 bit，相对 16 位少 16 倍），大约变成多久？

$2 \times 7 \times 0.1 + 2 \times 7/8 \times 0.65 \times 10^9 / (1.25 \times 10^8) ≈ 1.4 + 9.1 = 10.5$ 秒。压缩 16 倍后带宽项约 0.57 秒，加上 1.4 秒的延迟项约 2 秒——此时延迟（14 步 × 100 ms）成了大头。跨大洲时，减少通信的轮数（例如每隔几百步才同步一次的低通信方法，§8）和减少字节数同样重要。（注意：只发符号对 16 位梯度是 16 倍的压缩；代码 05 里以 fp32 为基准，是 32 倍。）

练习 11（概念：为什么 ZeRO-1 不增加通信） 普通数据并行的 all-reduce 和 ZeRO-1 的”reduce-scatter → 各自更新 → all-gather”通信量相同。那 ZeRO-1 是不是”免费的午餐”？有没有别的代价？

通信量相同，但有几处不同：(1) 普通数据并行可以把 all-reduce 拆成很多小桶、和反向传播重叠；ZeRO-1 的 all-gather 发生在优化器更新之后，要等更新完才能开始下一步的前向（实际实现会尽量和下一步的前向重叠）。(2) 每张卡只更新 $1/N_d$ 的参数，优化器的计算也省了。(3) 实现更复杂，检查点要合并各张卡上的优化器状态才能恢复到不同的卡数（BLOOM 和 OPT 都遇到过检查点格式转换的问题）。总的来说 ZeRO-1 非常接近免费，所以 BLOOM、Megatron 默认都开着它。

练习 12（概念：选择并行方式） 你有 64 张 A100 80GB（8 台机器 × 8 张，机器内 NVLink，机器之间 200 Gbit/s），要训练一个 30B 参数、上下文 4K 的模型。给出一个并行方案并说明理由。

一个合理的起点：(1) 模型状态 $16 \times 30\text{B} = 480$ GB，64 张卡用 ZeRO-3 / FSDP 每张卡 7.5 GB，或者张量并行 $t = 8$（机器内）+ ZeRO-1 数据并行 8 份。(2) 激活：4K 上下文、宽度约 7K 的 30B 模型，用 FlashAttention 后每层每序列约 $34sbh$ ≈ 1 GB（16 位），几十层就是几十 GB，需要选择性重算或张量 + 序列并行来切。(3) 一个具体方案：$t = 8$（机器内，切激活也切参数）× $d = 8$（跨机器的数据并行 + ZeRO-1），不用流水线（避免气泡）；如果激活仍然放不下，再加激活重算。另一个方案是纯 FSDP + 激活重算，实现简单，但每层都要 all-gather 参数，跨机器的 200 Gbit/s 可能成为瓶颈。最终要实测吞吐来选——就像 BLOOM 那样。

12. 延伸阅读 · Further reading

原课阅读（1/18 Model Training 以及 Syllabus “Training Methods and Infrastructure”）

• Shoeybi et al. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. 2019. arXiv:1909.08053
• Stas Bekman et al. The Technology Behind BLOOM Training. Hugging Face 博客，2022. huggingface.co/blog/bloom-megatron-deepspeed

Ce Zhang 的相关工作（§7 的依据）

• Yuan et al. Decentralized Training of Foundation Models in Heterogeneous Environments. NeurIPS 2022. arXiv:2206.01288
• Wang et al. Fine-tuning Language Models over Slow Networks using Activation Compression with Guarantees（AC-SGD）. NeurIPS 2022. arXiv:2206.01299
• Tang et al. 1-bit Adam: Communication Efficient Large-Scale Training with Adam’s Convergence Speed. ICML 2021. arXiv:2102.02888

补充

• Rajbhandari et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC 2020. arXiv:1910.02054
• Narayanan et al. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM. SC 2021. arXiv:2104.04473——3D 并行的经验法则、1F1B 与交错调度。
• Korthikanti et al. Reducing Activation Recomputation in Large Transformer Models. 2022. arXiv:2205.05198——激活公式、序列并行、选择性重算。
• Huang et al. GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. NeurIPS 2019. arXiv:1811.06965
• BigScience Workshop. BLOOM: A 176B-Parameter Open-Access Multilingual Language Model. 2022. arXiv:2211.05100
• Llama Team. The Llama 3 Herd of Models. 2024. arXiv:2407.21783——第 3.3 节的 4D 并行与可靠性。
• 本目录 CS336 第 7、8 讲（并行 1、2）：在真实的多 GPU 环境里实现这些并行方式。

13. 下一讲预告 · Next lecture

这一讲解决了”怎样把一次训练分给上千张卡”。但真正训练过大模型的人会告诉你，这只是开始：几千张卡连续跑两个月，每天都有卡坏掉；损失曲线会突然跳起来；16 位浮点数会溢出；你在训练中途不得不改学习率、换优化器、回滚到几天前的检查点。

第 10 讲：训练一个大模型——PaLM、OPT 与 BLOOM 的实录（1/25 客座 Aakanksha Chowdhery，3/1 客座 Susan Zhang）打开三本”训练日志”：PaLM 的架构选择（SwiGLU、并行层、多查询注意力、RoPE、去掉偏置）和它遇到的约 20 次损失尖峰；OPT-175B 团队公开的训练日志（最后的总结说一共记了 148 页笔记）里，两个月里的几十次手动重启、学习率调整和换优化器的尝试；BLOOM 为什么放弃 fp16。代码会在第 1 讲的模型上演示 fp16 的溢出与 loss scaling、人为制造训练不稳定并比较几种应对方法，以及用故障率估算”有效训练时间”。