第 13 讲 · ML 并行化（二）：零冗余与 PyTorch FSDP · Zero Redundancy and PyTorch FSDP

对应 Maps to：CMU 15-779（Fall 2025，Zhihao Jia），Lecture 13 · Fri 10/10 · ML Parallelization: Zero Redundancy and PyTorch Fully Sharded Data Parallelism。 必读论文（同第 12 讲）：S. Rajbhandari, J. Rasley, O. Ruwase, Y. He, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models（SC 2020）；S. Rajbhandari, O. Ruwase, J. Rasley, S. Smith, Y. He, ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning（SC 2021）。选读：J. Ren et al., ZeRO-Offload: Democratizing Billion-Scale Model Training（ATC 2021）。slides 引用的补充：Y. Zhao et al., PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel（VLDB 2023）。 依据说明：本讲以原课 slides 12-ML-parallelization-part1.pdf 的后半为主线：

• 大模型训练的挑战表；SGD 与 Adam 复习；
• 显存构成（slides 改编自 DeepSpeed 的报告：FP16 参数、FP16 梯度、FP32 的梯度 / 方差 / 动量 / 参数）；
• ZeRO-DP 的三个阶段（slides 逐步画了第 1 阶段：前向、反向、AllReduce、Adam 更新、all-gather FP16 参数）；
• 第 3 阶段”前向时广播参数、用完就丢、反向时再广播一次”；
• PyTorch FSDP：分成 FSDP 单元、分片、前向 all-gather、反向 all-gather + reduce-scatter，以及 fully_shard 的编程示例。

ZeRO 论文第 12 讲已读过”显存去哪儿了、数据并行的通信”那部分，本讲精读三个阶段、通信分析和 ZeRO-R；ZeRO-Infinity 在本讲精读。论文数字只做转述。

zero.py（模拟设备上的 ZeRO 1/2/3）、全部实验、ZeRO-Offload 划分的枚举、FSDP 时间线模拟器都是本讲义自己设计的；FlatParameter、包装粒度、预取、速率限制器、混合分片的细节来自 FSDP 论文，ZeRO-Offload / ZeRO-Infinity 的分析来自这两篇论文，都是 slides 之外的补充。 配套代码 Code：code/lecture13/（4 个脚本 + 共享模块 zero.py；只需 numpy；最长的约 5 秒） 预计阅读时间：3 小时（含运行代码）

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

0.1 学习目标

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

1. 说出混合精度 Adam 训练每个参数的 16 字节（或 20 字节）分别是什么，并解释为什么”放得下 FP16 权重”远远不等于”训得动”。
2. 讲清楚 ZeRO 的三个阶段各切掉了什么、每块卡还剩多少（4Ψ + 12Ψ/N、2Ψ + 14Ψ/N、16Ψ/N），并能复现论文表 1。
3. 推导每个阶段的通信量：前两个阶段和数据并行一样是 2Ψ，第三阶段 3Ψ；并解释为什么 slides 的第 1 阶段写法是 3Ψ。
4. 解释 ZeRO 为什么能和数据并行逐位相同，以及哪些实现细节（求和顺序、更新后的 all-gather）会破坏它。
5. 说清楚 PyTorch FSDP 的分片单元、FlatParameter、反向预取、速率限制器、混合分片分别解决什么问题。
6. 用 ZeRO-Offload 的数据流图推出”只把 FP16 参数留在 GPU”的唯一最优划分；用 ZeRO-Infinity 的 AIT 公式估算卸载到 CPU / NVMe 需要多少带宽。

0.2 先修与本讲范围

0.3 本讲地图

§1  大模型训练的显存挑战                              ← 代码 01（演示 1）
§2  显存都去哪儿了：模型状态                           ← 代码 01（演示 2–3）
§3  ZeRO-DP 的三个阶段                                 ← 代码 02（演示 1–2、5）
§4  通信分析：2Ψ 与 3Ψ                                 ← 代码 02（演示 3–4）
§5  PyTorch FSDP：把第 3 阶段做成工程                  ← 代码 03
§6  剩余状态：ZeRO-R                                   ← 代码 01（演示 4–5）
§7  把模型状态卸载出去：ZeRO-Offload 与 ZeRO-Infinity   ← 代码 04
§8  论文精读：ZeRO；ZeRO-Infinity；（补充）FSDP、ZeRO-Offload
§9 小结  §10 常见误区  §11 练习  §12 延伸阅读  §13 下一讲预告

0.4 关于 zero.py

zero.py 在”设备列表”上真的训练一个 5 层的多层感知机（Ψ = 32,656 个参数），4 块模拟设备：

• 混合精度：前向、反向用 np.float16，损失和 Adam 更新用 np.float32；
• 分片单元：每个 Linear 层的 W 和 b 拼成一个一维的 FlatParameter，补零到能被设备数整除，设备 d 负责第 d 段；
• 显存账本：每块设备上的每个数组都登记在 DeviceMemory 里（类别、字节数），随时知道当前用了多少、峰值是多少；
• 通信：用第 12 讲 collectives.py 的环形 reduce-scatter / all-gather，每一轮传输都记下来。

五种做法（数据并行、ZeRO-1 的两种写法、ZeRO-2、ZeRO-3）共用同一套前向 / 反向 / Adam 代码，只在”每块卡存什么、什么时候传、什么时候释放”上不同。这样做的好处是：结果是否相同、显存差多少、通信差多少，都是数出来的，不是算出来的。

1. 大模型训练的显存挑战 · The memory wall

slides 用一张表开场：

当时 V100 只有 16 / 32 GB，A100 40 / 80 GB。从 GPT-2 开始，一块卡就放不下了。

直觉：第 12 讲的数据并行让每块卡算一部分样本，计算被分摊了，存储却没有——每块卡都存一整份模型状态。卡越多，冗余越多：64 块卡就有 64 份一模一样的优化器状态。ZeRO（Zero Redundancy Optimizer）的想法就写在名字里：把冗余去掉。

为什么不直接用模型并行？ 模型并行（第 14 讲）把模型切开放到不同的卡上，确实能减少每块卡的存储，但：

• 它要求改写模型（切矩阵、插通信），只支持有限的算子；
• 层内切分的通信很频繁（每层都要通信），跨节点时效率急剧下降（第 12 讲引用的数字：40B 的 Megatron-LM 跨两个 DGX-2 节点，每块 V100 不到峰值的 5%）。

ZeRO 的目标是保留数据并行的简单和通信效率，同时得到模型并行的显存效率。

2. 显存都去哪儿了：模型状态 · Where does the memory go

2.1 Adam 复习

slides 复习了 SGD 和 Adam：每一步算出梯度 g，SGD 直接 w ← w − ηg；Adam 还要维护两个”状态”：

• 一阶矩（动量）m ← β₁m + (1 − β₁)g；
• 二阶矩（方差）v ← β₂v + (1 − β₂)g²；
• 更新 w ← w − η · m̂ / (√v̂ + ε)（m̂、v̂ 是偏差修正后的值）。

这两个状态和参数一样大，而且只在更新那一刻才用到。

2.2 混合精度：为什么要 FP32 的主参数

混合精度训练用 FP16（或 BF16）做前向和反向，因为 Tensor Core 在低精度下快得多。但更新必须在 FP32 里做：学习率乘梯度常常比参数小 3–4 个数量级，FP16 只有 10 位尾数，w + Δw 会直接舍入回 w——更新被”吃掉”了。所以优化器保留一份 FP32 的主参数（master weights），每步在 FP32 上更新，再转成 FP16 给下一步用。

2.3 每个参数多少字节（代码 01 演示 1）

三种说法的差别只在 FP32 梯度算不算常驻状态：

• ZeRO 论文用 K 表示优化器状态的字节数，混合精度 Adam 的 K = 12（主参数、两个矩），合计 2Ψ + 2Ψ + KΨ = 16Ψ；
• slides（改编自 DeepSpeed 的报告）和 ZeRO-Infinity 把 FP32 梯度也算进去，合计 20Ψ，于是 1B 模型是 slides 上的”20GB/GPU”。

代码里用 zero.py 跑一步数据并行，反向刚结束时设备 0 上恰好是 16.00 字节 / 参数；FP32 梯度是更新时一次一个单元转出来的临时量，峰值只多出最大单元的 4 字节 / 参数。

为什么这很重要：GPT-2 1.5B 的 FP16 参数只有 3 GB，但训练要 16 × 1.5 = 24 GB——八倍。“放得下权重”和”训得动”是两回事。而 16 字节里有 12 字节是只在更新时才用到的优化器状态。

2.4 复现 ZeRO 表 1（代码 01 演示 2）

ZeRO 论文图 1 给出了三个阶段每块卡的模型状态：

代码按公式算出 7.5B、128B、1T 三个模型在 N_d = 1 到 1024 下的 54 个数，和论文表 1 逐一比较，最大相对差 2.34%，全部来自论文的取整（0.117 写成 0.12，4187.5 写成 4187）。几个值得记住的数：

• 7.5B 模型、64 卡：数据并行 120 GB，P_os 31.4 GB，P_os+g 16.6 GB，P_os+g+p 1.88 GB；
• 1T 模型、1024 卡、第 3 阶段：每卡 15.6 GB——放得进 32 GB 的 V100。

前两个阶段有一个”地板”：不管多少卡，每块卡都还要存一整份 FP16 参数（第 1 阶段还有梯度），所以 7.5B 模型在第 1 阶段永远不低于 30 GB。只有第 3 阶段随卡数线性下降。

2.5 加上模型并行（代码 01 演示 3）

ZeRO 论文表 2 问：MP 路模型并行 × N_d 路数据并行，32 GB 的卡最大能训多大的模型？每块卡只放 Ψ/MP 个参数，再在 N_d 路里用 ZeRO 切，最大模型 = 32 GB ÷ (每参数字节数) × MP。代码算出的理论值和论文完全一致，例如 MP = 1、64 卡：baseline 2B、P_os 7.6B、P_os+g 14.4B、P_os+g+p 128B；MP = 16、1024 卡：P_os+g+p 2T。

论文还给出了实测值（只开 P_os）：MP = 1 时 6.2B，比理论的 7.6B 小——因为真实训练还要放激活、临时缓冲区，还有碎片（§6）。

3. ZeRO-DP 的三个阶段 · The three stages

3.1 第 1 阶段：切优化器状态（P_os）

slides 一步一步地画了第 1 阶段（两块 GPU，每块卡的显存画成按层排列的格子）：

1. 切分：把 FP32 的主参数、动量、方差均分成 N 份，GPU i 只存第 i 份；
2. 前向：和数据并行一样，每块卡用完整的 FP16 参数算自己的样本；
3. 反向：每块卡得到完整的 FP16 梯度，然后 AllReduce 求平均；
4. 更新：GPU i 只用 Adam 更新自己负责的那 1/N（只需要梯度的第 i 段）；
5. 更新 FP16 参数：把更新后的 FP32 主参数转成 FP16——但每块卡只更新了自己那 1/N；
6. all-gather FP16 参数：每块卡把自己那段发给大家，所有卡都拿到完整的新参数，这一步迭代结束。

为什么可以这样切：Adam 是逐元素的——第 j 个参数的更新只用到第 j 个梯度、第 j 个 m 和 v。所以把参数切成 N 段、各自更新，和一起更新在每个元素上做的运算完全相同。

3.2 第 2 阶段：再切梯度（P_os+g）

既然 GPU i 只用梯度的第 i 段，其余 (N − 1)/N 的梯度就不必保留。做法：

• 反向时每一层的梯度一算完，就把它 reduce 到”负责这些参数的那块卡”上，然后立刻释放；
• 实际中按桶（bucket）做，相当于对每个桶做一次 reduce-scatter，和 DDP 的”梯度分桶 + 通信重叠”同一个思路（第 12 讲 §6.2），只是 AllReduce 换成了 reduce-scatter。

slides 的写法：“Perform AllReduce right after back propagation of each layer; only one GPU keeps the gradients after AllReduce”——说的就是这件事。

第 1 阶段和第 2 阶段的显存区别在哪里：第 1 阶段在整个反向期间都持有完整的梯度缓冲区，反向结束才处理；第 2 阶段在反向进行中就把不归自己管的梯度释放了。

3.3 第 3 阶段：连参数也切（P_os+g+p）

slides 用两个卷积层的例子：数据并行时每块卡都存 W1、W2；第 3 阶段 GPU1 只存 W1，GPU2 只存 W2。

• 前向：算到某一层之前，负责它的卡把参数广播给所有卡；算完立刻丢掉；
• 反向：反向算到这一层时，再广播一次；
• 梯度照第 2 阶段那样 reduce-scatter。

这样每块卡长期只存 1/N 的一切，同时拼完整的只有正在算的那一两层。论文的实现（以及后来的 FSDP）用 all-gather 而不是广播：每个参数被切给所有卡，用的时候各卡把自己那段拿出来拼。通信量相同（§4），但 all-gather 在卸载到 CPU 时有大优势（§7.3）。

3.4 代码验证：逐位相同（代码 02 演示 1）

4 块卡，每卡每步 16 个样本，训练 50 步：

ZeRO 不是近似，它和数据并行做的是完全相同的计算。逐位相同依赖三件事：

1. 前向、反向用的 FP16 参数在每块卡上都完全相同；
2. 梯度的求和顺序相同。浮点加法不满足结合律，(a + b) + c 和 a + (b + c) 的舍入可能不同。环形 AllReduce 本身就是 reduce-scatter + all-gather，只要两边用同一个环、同样的分段，每个元素的加法顺序就一样。为此代码里的数据并行也按单元（像 DDP 的桶）做 AllReduce；如果它对整个梯度向量做一次 AllReduce，分段不同，结果就只是”很接近”；
3. Adam 是逐元素的。

3.5 代码验证：显存（代码 02 演示 2）

反向刚结束时设备 0 上的数组（字节）：

每一格都和公式逐字节相等（N = 4：16 → 4 + 12/4 = 7 → 2 + 14/4 = 5.5 → 16/4 = 4）。

峰值 > 稳态：ZeRO-3 的峰值出现在反向的某个单元上——此刻设备 0 上除了常驻的 1/N 分片，还有这个单元 all-gather 出来的完整 FP16 参数（p16_full，18,624 字节）和刚算出、还没 reduce-scatter 的完整梯度（g16_full，18,624 字节）。FSDP 论文把峰值写成 O(Ψ/F + max_i ψ_i)：常驻的分片 + 最大的那个单元。

3.6 设备数变化与补零（代码 02 演示 5）

• N = 1 时四种做法完全一样：ZeRO 的收益全部来自”有 N 份冗余可以消除”。
• N = 3 时，最后一层的 1,552 个参数不能被 3 整除，要补 2 个零。补的零也要存、也要传，公式里的 Ψ 要换成 Ψ’。

4. 通信分析：2Ψ 与 3Ψ · Communication volume

4.1 论文的分析

ZeRO 论文 §7 按”每块卡收发的元素数”计算通信量：

• 数据并行：AllReduce = reduce-scatter（Ψ）+ all-gather（Ψ）= 2Ψ；
• P_os、P_os+g：梯度只需要 reduce-scatter（Ψ），更新后 all-gather 参数（Ψ），合计 2Ψ——和数据并行一样。ZeRO 只是把 AllReduce 后半段的 all-gather 的对象，从”求和后的梯度”换成了”更新后的参数”；
• P_os+g+p：前向 all-gather 参数（Ψ）、反向再 all-gather 一次（Ψ）、梯度 reduce-scatter（Ψ），更新后不用再 all-gather（每块卡本来就只存自己那段），合计 3Ψ = 1.5 倍。

4.2 代码验证（代码 02 演示 3）

设备 0 每一步发出的字节（以 2Ψ 字节 = Ψ 个 FP16 元素为单位；环形算法里每块卡发出 (N − 1)/N 份，所以 2Ψ 对应 1.50、3Ψ 对应 2.25）：

一个容易忽略的细节：slides 写的第 1 阶段（先 AllReduce 梯度、再 all-gather 参数）比论文多传了一份——AllReduce 后半段把求和后的梯度 all-gather 给所有卡，但每块卡只用自己那段，其余的白传了。两种写法显存相同（反向期间都持有完整梯度），通信差 1.5 倍。

ZeRO-3 不仅字节多，轮数也多：每个单元 3 次集合通信，每次 N − 1 轮。消息小的时候延迟 α 会更明显，这是 FSDP 要把参数拼成 FlatParameter、按单元通信的原因（§5.2）。

4.3 一个常见错误（代码 02 演示 4）

ZeRO-2 更新完自己那段以后忘了 all-gather FP16 参数：每块卡只把自己负责的 1/4 更新了，其余 3/4 一直停在初始值。各卡的”副本”从此不同（最大差 0.08–0.09），损失 1.68（正确的是 0.89）。

这个错误不会报错、也不会崩溃。实践中常用的检查：定期比较各卡参数的校验和或范数，它们必须完全相等。

5. PyTorch FSDP：把第 3 阶段做成工程 · Fully Sharded Data Parallel

5.1 slides 的四步

slides 把 FSDP 概括为：

1. 把模型参数分成若干 FSDP 单元（units）；
2. 把每个单元切分到多块 GPU 上；
3. 前向：对每个单元先 all-gather 拿到完整参数，算完就丢掉别人的分片；
4. 反向：再 all-gather 一次；每块卡算出这个单元全部参数的梯度；然后 reduce-scatter 聚合梯度，每块卡只留自己那段。

和 ZeRO-3 完全对应。FSDP 论文（VLDB 2023）讲的是把它做进 PyTorch 时遇到的工程问题。

5.2 FlatParameter：拼接、补零、均分（代码 03 演示 1）

一个单元里可能有几十个参数张量（权重、偏置、LayerNorm……）。FSDP 把它们展平、首尾相接成一个一维的 FlatParameter，右边补零到能被分片数 F 整除，再均分成 F 段，每块卡一段。FSDP 论文图 3 的例子：一个 4×3 的 nn.Linear（12 个权重 + 3 个偏置 = 15 个元素）分到 16 块卡，补 1 个零，每块卡 1 个元素，最后一块卡拿的是补的零。

这样做有两个好处：

• NCCL 的 all-gather / reduce-scatter 要求各卡输入等长，而且输出要写到一块连续的张量里。FlatParameter 的布局正好是集合通信要的布局，通信前后都不用额外拷贝。FSDP 论文实测过：不等长时 PyTorch 只能用一组广播来模拟，明显更慢；
• 一个单元一次通信：通信次数 = 单元数，而不是参数张量的个数。

补零最多 F − 1 个元素，对真实模型可以忽略（一个 3.1 亿参数的块分到 1000 块卡只补 200 个）。

5.3 单元切多大（代码 03 演示 2）

FSDP 论文说：峰值参数显存是 O(Σψᵢ/F + max ψᵢ)，集合通信次数是 O(单元数)。单元越小，同时拼完整的参数越少，但通信次数越多。13B 模型、64 卡全分片的估算：

“每个 Transformer 块一个单元”是常见的折中，slides 的编程示例正是这样写的：

from torch.distributed.fsdp import fully_shard, FSDPModule
from torch.distributed.tensor import DTensor

model = Transformer()
for layer in model.layers:
    fully_shard(layer)          # 每一层一个 FSDP 单元
fully_shard(model)              # 剩下的参数（嵌入层等）归最外层的单元

for param in model.parameters():
    assert isinstance(param, DTensor)       # 分片后参数变成 DTensor

optim = torch.optim.Adam(model.parameters(), lr=1e-2)   # 优化器自动只管自己的分片
for _ in range(epochs):
    x = torch.randint(0, vocab_size, (batch_size, seq_len))
    loss = model(x).sum()
    loss.backward()
    optim.step()
    optim.zero_grad()

训练循环和单卡一模一样：all-gather、释放、reduce-scatter 都藏在 fully_shard 装的前向 / 反向 hook 里。

FSDP 论文图 2(b) 还测了”一次 all-gather 多大才划算”：总量固定为 2³⁰ 个 FP32 元素，每次的元素数少于约 3,300 万后，总时间迅速增加。代码的 α-β 模型只有 α = 5 µs 一项固定开销，3,300 万时只慢 13%，趋势相同但平缓得多——真实的 NCCL 调用还有 kernel 启动、同步等开销。

5.4 通信和计算重叠（代码 03 演示 3）

FSDP 用一条单独的 CUDA 流做通信。为什么要单独的流？ProcessGroupNCCL 在发起集合通信前会和当前流同步。DDP 的反向是”先算、后通信”，没问题；FSDP 的前向却是”先通信（all-gather）、后算”，而且在 eager 模式下 FSDP 事先不知道下一个要 all-gather 哪个单元。如果 all-gather 发在计算流上，它就要等前面的计算做完，重叠不起来。

代码写了一个两条流的时间线模拟器：CPU 按程序顺序发射操作，每条流先进先出，一个操作的开始时间 = max(发射时间, 这条流空闲, 依赖的操作结束)。4 个单元的小例子（F 前向、B 反向、A all-gather、R reduce-scatter）：

不重叠（总时间 40.5，纯计算 24.0）
  计算 |   F0FF   F1FF   F2FF   F3FFB3BBBBBB      B2BBBBBB      B1BBBBBB      B0|
  通信 |A0A    A1A    A2A    A3A            R3RA2A        R2RA1A        R1RA0A  |
重叠，不做反向预取（36.0）
  计算 |   F0FFF1FFF2FFF3FFB3BBBBBB      B2BBBBBB      B1BBBBBB      B0BBBBBB   |
  通信 |A0AA1A A2A A3A             R3RA2A        R2RA1A        R1RA0A        R0R|
重叠 + 反向预取（27.0）
  计算 |   F0FFF1FFF2FFF3FFB3BBBBBBB2BBBBBBB1BBBBBBB0BBBBBB                     |
  通信 |A0AA1A A2A A3A A2A         R3RA1A  R2RA0A  R1R     R0R                  |

反向预取（backward prefetch）：FSDP 每块卡只有一个进程组，all-gather 和 reduce-scatter 在同一条 NCCL 流上排队。不预取时，反向的发射顺序是 RS(i)（单元 i 的反向一结束）→ AG(i − 1)（单元 i − 1 的反向开始前）。BWD(i − 1) 要等 AG(i − 1)，AG(i − 1) 又排在 RS(i) 后面，每个单元都暴露两次通信（图里 B 之间的空白）。反向预取把 AG(i − 1) 提到 RS(i) 之前发出，让它和 BWD(i) 重叠。FSDP 用”上一次前向的逆序”作为反向顺序的代理来知道下一个是谁。

13B 模型、64 卡全分片、每卡 16,384 token 的估算（每块前向 20.8 ms、反向 41.7 ms、一次 all-gather / reduce-scatter 13.0 ms）：

在这个模型里反向预取快 40%。FSDP 论文在 GPT-175B 上测得约 18%——数字不可比（模型、集群都不同），趋势一致。

速率限制器（rate limiter）：PyTorch 的 CUDA 缓存分配器在 CPU 发射操作的那一刻就要决定用哪块显存，而不是等 GPU 真的执行时。CPU 通常跑得比 GPU 快：不限速的话，CPU 一口气把前向的 40 个和反向的 39 个 all-gather 全发出去，每个的目标缓冲区都立刻分配，峰值是 79 个块的缓冲区——比整个模型的 FP16 参数还多一倍，等于没有分片。更糟的是，分配给通信流的块不能被计算流复用，可能逼得分配器做 cudaFree / cudaMalloc 重试（会同步整个设备）。FSDP 的速率限制器最多允许 2 个在途的 all-gather（还能重叠的最小值）。代码里限速后时间完全不变，显存从 79 块降到 2 块。FSDP 论文实测限速器的效果”并不一致”：T5 上有明显加速，RegNet 上没有，DeepViT 上反而变慢——它只在 CPU 跑得太快、真的触发了分配器重试时才有用。

前向预取（forward prefetch）：如果 CPU 慢（例如 Python 开销大），all-gather 要等 CPU 走到下一个单元才发出。前向预取在当前单元计算之前就发下一个单元的 all-gather。代码的模型里它几乎没有效果（< 0.5%）：这里把一个单元的计算当成一个操作，CPU 受限时 all-gather 在 CPU 的空闲里本来就做完了。论文的场景是 CPU 发不出足够多的 all-gather 来”填满 NCCL 流”，要把一个单元拆成许多 kernel、逐个建模发射时间才看得出来。这是这个模拟器的局限。

5.5 混合分片（代码 03 演示 4）

FSDP 引入分片因子 F：参数在 F 块卡之间分片。F = 1 就是 DDP，F = 总卡数 W 是全分片，中间是混合分片（hybrid sharding）：在 F 块卡的”分片组”内切分，在 W/F 个分片组之间复制。

梯度为什么还对？全局的和可以分两步加：Σ_所有卡 g = Σ_各分片组 (Σ_组内 g)。先在组内 reduce-scatter，每块卡得到本组某一段的和；再在”各组同一位置的卡”组成的复制组里 AllReduce 这一段（FSDP 论文式 (1)）。代码在 2 节点 × 8 卡上真的做了一遍，和直接求和差 2.0e-15。

最自然的选法是 F = 每个节点的卡数：all-gather 和 reduce-scatter 都在节点内的 NVLink 上，只有梯度分片的 AllReduce 跨节点。13B 模型、8 节点 × 8 卡的估算：

全分片显存最省，但三次跨节点的环形集合通信要 1.51 s，已经和计算（2.53 s）差不多。混合分片把跨节点流量降到约 1/13，代价是每卡多存一些模型状态。FSDP 论文还指出，在较小的组里做集合通信受”掉队者”（straggler）和网络干扰的影响也更小。F 让显存和跨节点流量之间可以连续地取舍，这对”全复制刚好放不下、全分片又浪费”的中等规模模型尤其有用。

5.6 其他工程细节（FSDP 论文）

• 延迟初始化（deferred initialization）：模型大到一块卡放不下时，连构造模型都做不到。FSDP 先在一个”假设备”（meta device）上构造模型，只记录初始化操作；然后一次把一个单元搬到 GPU、重放它的初始化、立刻分片，再处理下一个。
• 和 autograd 配合：前向前把原始参数设成 FlatParameter 的视图（用 autograd 可见的 split / view），autograd 自然地把梯度写到 FlatParameter 梯度的对应位置；在 FlatParameter 的梯度完成时触发一个 hook 做 reduce-scatter。
• 梯度累积：可以每个 micro-batch 都 reduce-scatter（省显存），也可以先在本地累积完整的梯度、最后再通信（省通信、多用显存）。

论文的结果：T5-611M、2.28B 时 FSDP 和 DDP 性能相近，更大的 T5 DDP 就 OOM 了；GPT-175B 在 128 到 512 块 A100 上近似线性扩展，每卡 173 / 186 TFLOPS（batch 1 / 2，约为 BF16 峰值的 55% / 60%）。

6. 剩余状态：ZeRO-R · Residual states

ZeRO-DP 只处理模型状态。模型状态切小以后，剩下的三类就成了新的瓶颈。

6.1 激活（代码 01 演示 4a）

ZeRO 论文脚注给出 GPT-2 式结构的全部激活 ≈ 12 × hd × batch × seq × 层数 个元素。代码按这个公式和 ZeRO-Infinity 的检查点公式（每块存一个，2 × bsz × seq × hd × nl / ci 字节）计算：

一处对不上的数字：ZeRO §6.1 说同一个 100B 模型、MP = 16 时每块卡存检查点”约 33 GB”，分片后”约 2 GB”。按公式是 67.1 GB 和 4.2 GB，正好差一倍。论文没给出这个数的推导，可能用了不同的 batch 假设。这里如实记下，不去凑。

P_a（激活分片）：模型并行的 MP 块卡算的是同一批样本，每块卡都存一份完整的激活检查点——又是冗余。ZeRO-R 把检查点切成 MP 份，反向重算前再 all-gather。通信开销：Megatron-LM 每个 Transformer 块前向 2 次、重算 2 次、反向 2 次 AllReduce，共 12 × seq × hd；P_a 每块只多一次 all-gather（seq × hd），不到 10%。显存省 MP 倍，于是可以把 batch 放大，反过来减少数据并行的通信。极端情况还可以把分片后的检查点卸载到 CPU（P_a+cpu）。

6.2 临时缓冲区（代码 01 演示 4b）

做 AllReduce 或求梯度范数时，常把所有梯度拼成一个大的 FP32 缓冲区（大消息带宽更高）。1.5B 模型要 6 GB，3B 要 12 GB，175B 要 700 GB——缓冲区和模型一样大。C_B（固定大小的缓冲区）：缓冲区大小不随模型增长，只要足够大、通信效率够高就行，和 DDP 的 25 MB 桶是同一个想法。

6.3 碎片（代码 01 演示 5）

长命的对象（激活检查点、参数梯度）和短命的对象（被丢掉的激活、激活梯度）交替分配，会把空闲空间切碎。代码用一个首次适配分配器模拟：9 层前向，每层分配大小不一的临时激活、存一个检查点、释放临时激活。结束时空闲 44%，最大的连续空洞只有 260，一个 400 的请求就失败了。M_D（碎片整理）：把检查点和参数梯度复制到预先分配的连续缓冲区里，空闲空间保持连续，同样的请求成功。论文说极端情况下”还有 30% 以上空闲时就 OOM”。FSDP 论文里 128 卡、batch 2 的 GPT-175B 每卡 TFLOPS 明显偏低，也是因为碎片整理：反向占到迭代时间的 85.56%，而正常约 67%。

6.4 单个算子的工作内存

ZeRO-Infinity 定义了 MSWM（Model State Working Memory）：即使模型状态都卸载出去，算一个算子时它的参数和梯度也必须在 GPU 上，而且是连续的。最大的 Linear（hd → 4hd）的 FP16 参数 + 梯度是 4 × hd × 4hd 字节：hd = 8K 时 1.1 GB，hd = 64K 时 68.7 GB——一个算子就放不进一块卡。ZeRO-Infinity 的以内存为中心的分块（memory-centric tiling）把一个大 Linear 拆成几个小 Linear 顺序执行，各自的参数单独 all-gather / 释放。论文实验：人为把显存切成 2 GB 的块后，不分块最大只能训 hidden 8K，分块因子 16 时能训 64K。

7. 把模型状态卸载出去 · Offloading to CPU and NVMe

7.1 ZeRO-Offload：唯一最优的划分（代码 04 演示 1）

ZeRO-Offload 把训练画成一张数据流图：数据节点是 p16、g16、p32、m32、v32，计算节点是 FWD、BWD、更新（update）、FP32 → FP16 转换（float2half）。边的权重是一次迭代流过的数据量：FP16 状态 2M，FP32 状态 4M（M 是参数个数）。任何 GPU / CPU 划分的 PCIe 通信量 = 被切断的边的权重之和。论文用三条原则筛选：

1. CPU 不做 O(MB) 的计算：前向反向的计算量和 batch B 成正比，必须留在 GPU，所以 FWD 和 BWD 合成一个超级节点放在 GPU；只有 O(M) 的计算（更新、范数）可以卸载；
2. 通信最少：
   • FWD-BWD 固定后图里每个节点都在一个环上，任何卸载都至少切断两条边，至少 4M；
   • 要达到 4M，FP32 的三个状态必须和它们的生产者、消费者（update、float2half）在同一边，否则要切一条 4M 的边，至少 6M；
   • p16 必须和 FWD-BWD 在一起（它们之间是 4M）；
3. 显存最省：剩下 g16 和”更新超级节点”两个自由度，正好是论文表 1 的四行。

代码把 FWD-BWD 固定在 GPU，枚举其余 7 个节点的全部 128 种划分。通信 ≤ 4M 的恰好是表 1 的四行：

在通信 4M 的划分里显存最省的只有一个：GPU 上只留 FP16 参数，FP16 梯度、全部 FP32 状态和 Adam 都在 CPU。论文说它”唯一最优”，在这张图上可以直接枚举验证。

顺便一提：枚举里通信 6M 的那一行有一种划分让 GPU 显存为 0——连 FP16 参数也放到 CPU，每次用时再搬上来。ZeRO-Offload 不选它（多 2M 的 PCIe 通信），而 ZeRO-Infinity 恰恰选了它，再用 §7.3 的”以带宽为中心的分片”把多出来的通信扛住。

7.2 ZeRO-Offload 的时间线与 DPU（代码 04 演示 2）

单块 GPU 的一步：

• GPU 前向（FP16 参数已在 GPU 上，不用通信）；
• 反向时每个梯度一算完就经 PCIe 下行到 CPU（和反向重叠）；
• CPU 上的 Adam 更新 FP32 状态；
• 新的 FP16 参数分块上行到 GPU（和 CPU-Adam 的下一块流水）。

两个让它不慢的技术：

• CPU-Adam：用 SIMD、多线程、分块拷贝写的 CPU Adam。论文表 4：10B 参数 2.57 s，PyTorch 的 CPU Adam 14.76 s，GPU 上 1.00 s；
• 一步延迟的参数更新（DPU）：batch 小的时候 GPU 计算并不比 CPU-Adam 长多少。DPU 让第 i 步的 CPU 更新和第 i + 1 步的 GPU 前向反向同时进行，代价是第 i + 1 步用的参数用的是第 i − 1 步的梯度。论文先正常训练几十步再打开 DPU，实测不影响收敛。

代码的估算（10B 参数，V100 FP16 峰值的 50%，有激活检查点）：

论文：10B 模型在单块 V100 上超过 40 TFLOPS；DPU 在 micro batch 8 时带来 1.12–1.59 倍（不同模型大小）。估算给出的是同一量级的趋势，不是复现。

7.3 ZeRO-Infinity：要多少带宽（代码 04 演示 3）

ZeRO-Infinity 先问：卸载到 CPU / NVMe，带宽够不够？它定义算术强度 AIT = 计算量 / 需要搬的数据量，不重叠时：

$\text{效率} = \frac{\text{AIT} \cdot bw}{\text{AIT} \cdot bw + \text{peak}}$

Transformer 每步的计算 ≈ 2 × 4 × 12 × bsz × seq × nl × hd²（前向 2、反向 4、重算 2，乘参数数 12·nl·hd²）。三类数据的 AIT：

• 参数和梯度：参数要搬 3 次（前向、反向、重算），梯度写 1 次，共 4 × 2 字节 × 参数数 → AIT = seq × bsz；
• 优化器状态：读写各一次 × 16 字节 → AIT = seq × bsz / 4；
• 激活检查点：存取各一次 × 2 × bsz × seq × hd × nl / ci 字节 → AIT = 24 × hd × ci。

取峰值 70 TFLOPS、序列长度 1024，代码复现了论文 §5.2 的结论：

结论：参数 / 梯度要 70 GB/s 以上，优化器状态要 1.5 TB/s 以上，激活检查点 1–4 GB/s 就够。

• 优化器状态要求最高，但它的更新可以在所有设备上并行，带宽是所有 GPU / CPU 加起来的；
• 激活检查点的 AIT 和 hidden 成正比，大模型可以放心卸载；
• 真正的难点是参数和梯度：一条 PCIe 只有约 12 GB/s。

7.4 以带宽为中心的分片（代码 04 演示 4）

ZeRO 和 ZeRO-Offload 里每个参数属于一块卡（owner），要用时 owner 广播给大家。数据在 GPU 上时广播和 all-gather 通信量一样；但数据在 CPU / NVMe 上时，广播要先经 owner 那一条 PCIe 搬上 GPU，其他 N − 1 条 PCIe 闲着。ZeRO-Infinity 把每个参数都切给所有卡，用的时候 N 块卡各经自己的 PCIe 取 1/N，再在 GPU 之间 all-gather——有效带宽随卡数线性增长：

（每卡 3.0 / 1.6 GB/s 是论文给出的”所有 GPU 同时读时”的带宽；有效带宽封顶于 GPU 之间的 70 GB/s。论文：64 个 DGX-2 节点上 CPU 内存聚合带宽超过 3 TB/s，NVMe 超过 1.5 TB/s。）

另外两个效率技术：

• 以重叠为中心的设计：从 NVMe 取参数分三步（NVMe → CPU、CPU → GPU、GPU 之间 all-gather），动态预取器记录每次迭代的算子顺序，把下一个参数的三步和当前计算重叠；
• infinity offload engine：接近峰值的 NVMe 读写，读写和 CPU 计算重叠，复用临时缓冲区以减少碎片。

7.5 能训多大（代码 04 演示 5）

（估算给激活和工作内存留了 4 GB，ZeRO-Infinity 按 20 字节 / 参数。）

ZeRO-Infinity 让模型大小只受 CPU 内存和 NVMe 容量限制。能不能训得快，要靠 §7.3、§7.4 的带宽分析。

8. 论文精读 · Paper reading

8.1 ZeRO: Memory Optimizations Toward Training Trillion Parameter Models（Rajbhandari et al., SC 2020）——三个阶段与 ZeRO-R

1. 问题：数据并行计算和通信效率高、好用，但每块卡存一整份模型状态；模型并行省显存，但难用、跨节点效率差。能不能两者兼得？

2. 为什么重要又困难：

• 模型规模每年翻好几倍，显存只缓慢增长；
• 消除冗余看起来必然要多通信（“要用别人的参数，就得把它传过来”）；
• 剩余状态（激活、缓冲区、碎片）在模型状态切小以后会成为新的瓶颈。

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

• 数据并行冗余 N 份；
• 模型并行（Megatron-LM）需要改模型、跨节点慢；
• 流水线并行（GPipe、PipeDream）有空泡或权重暂存的显存开销（第 14 讲）；
• CPU 卸载（L2L 等）受 PCIe 带宽限制；
• 激活检查点、压缩只处理激活。

4. 新意：

• 把显存分成模型状态和剩余状态两类，分别处理；
• ZeRO-DP：利用”并不是所有状态都在所有时刻被需要”（优化器状态只在更新时用，某一层的参数只在算这一层时用），用动态的通信调度把状态切开，并证明前两个阶段通信量不变、第三阶段只多 50%；
• ZeRO-R：激活分片（P_a）、固定大小的缓冲区（C_B）、碎片整理（M_D）；
• ZeRO 和模型并行可以组合，理论上每卡显存降为 1/(N_d × N_m)。

5. 结果：

• ZeRO-100B（P_os+g + ZeRO-R）在 400 块 V100 上训练 170B 参数的模型，8B–100B 时持续约 15 PFLOP/s（超过峰值的 30%）；
• 比当时最好的方法模型大 8 倍、快最多 10 倍；60B 模型从 64 卡到 400 卡超线性扩展（卡多了每卡显存少了，batch 能变大）；
• 不用模型并行就能训 13B；Turing-NLG（17B）用它训练。

本讲的验证：在 4 块模拟设备上实现了三个阶段，结果和数据并行逐位相同，每卡显存和公式逐字节相等，通信量 2Ψ / 3Ψ 由通信记录数出。

局限：

• 第 3 阶段的通信（每层前向、反向各一次 all-gather）对网络要求高，跨节点时很难完全藏住（§5.5）；
• 论文发表时只实现了 P_os+g，第 3 阶段后来在 DeepSpeed 和 FSDP 里落地；
• 不减少计算，也不减少全局 batch 随卡数增长的问题。

8.2 ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning（Rajbhandari et al., SC 2021）

1. 问题：三年里最大的稠密模型增长了 1000 多倍，单卡显存只增长了 5 倍（16 → 80 GB）。即使用 3D 并行，光是放下一个万亿参数的模型就要 800 块 V100（80 GB 的 A100 也要 320 块）；微调 GPT-3 也要 8 个 DGX-2 节点。能不能用 GPU 之外的存储？能不能不改模型代码？

2. 为什么困难：CPU 内存带宽比 GPU 显存低一个数量级，NVMe 再低一个数量级，从 CPU / NVMe 到单块 GPU 还要过约 12 GB/s 的 PCIe。直觉上，卸载必然很慢。

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

• 3D 并行要改模型、切算子，而且受总显存限制；
• ZeRO-Offload 只卸载优化器状态和梯度，FP16 参数仍在 GPU 上（单卡最多约 13B），而且参数经 owner 一条 PCIe 搬运；
• 早期的 NVMe 方案针对推荐模型的稀疏参数。

4. 新意：

• 先算账：给出模型状态、激活、两种工作内存（MSWM、AWM）的显存公式，以及参数 / 梯度、优化器状态、激活检查点的 AIT 和带宽需求（70 GB/s、1.5 TB/s、1–4 GB/s）——它把”卸载能不能快”变成了一个可以计算的问题；
• infinity offload engine：把所有模型状态卸载到 CPU 或 NVMe；
• memory-centric tiling：把超大的单个算子拆成小块顺序执行，不用模型并行也能训超宽的层；
• 以带宽为中心的分片：每个参数切给所有卡、all-gather 代替广播，让所有 PCIe 一起工作；
• 以重叠为中心的设计：动态预取器，把 NVMe → CPU、CPU → GPU、GPU 之间的三步传输和计算重叠；
• 以易用为中心的实现：在 PyTorch 子模块上自动插前向 / 反向 hook 做 all-gather 和分片，包装构造函数让模型在初始化时就分片，整个模型从不在一个进程里完整出现。

5. 结果（512 块 V100、32 个 DGX-2 节点）：

• 训练 32 万亿参数的模型，是 3D 并行（约 650B）的 50 倍；
• 20 万亿参数时仍有 34 TFLOPS / GPU，10 万亿时 43 TFLOPS / GPU；整体超过 25 PFLOP/s（峰值的 40%）；
• 1T 模型从 4 节点到 32 节点超线性扩展（聚合 PCIe / NVMe 带宽随节点增长）；
• 单个 DGX-2 节点：100B 以下超过 40 TFLOPS / GPU，最大能训 1T 参数，比单纯的数据并行大 700 倍；
• 以带宽为中心的分片让梯度卸载在 64 卡时比 ZeRO-Offload 快近 2 倍。

局限：

• 20T 时性能下降是因为 CPU 内存放不下更大 batch 的激活检查点，batch 极小；
• 单节点、小 batch 时带宽仍是瓶颈（这正是 AIT 分析预言的）；
• 论文最后一张表提醒：如果加速器再快 10 倍、100 倍，需要的慢速内存带宽和 GPU 间带宽也要相应增长 10 倍、100 倍——卸载的”免费”是有条件的。

8.3 两篇论文之间的联系

• 同一个记账框架：ZeRO 回答”每块卡要存多少”，ZeRO-Infinity 在此基础上回答”存到更慢的地方，需要多少带宽”。两篇都先用公式算清楚，再设计系统。
• 同一个洞察，用到不同的层次：“不是所有状态都在所有时刻被需要”。ZeRO 用它在卡之间切分（用时再 all-gather），ZeRO-Infinity 用它在存储层次之间搬运（用时再从 CPU / NVMe 取）。
• 广播 vs all-gather：ZeRO 论文第 3 阶段描述成”owner 广播”，ZeRO-Infinity 改成”每个参数切给所有卡、all-gather”。在 GPU 上两者通信量一样；在 CPU / NVMe 上，后者让所有 PCIe 一起工作。一个看起来无关紧要的选择，换了一个场景就成了关键。
• 一个可以讨论的问题：ZeRO 的超线性扩展来自”卡多了每卡显存少了，batch 能变大”；ZeRO-Infinity 的超线性扩展来自”节点多了聚合带宽变大”。两者都说明：在显存受限的区域，“更多的卡”带来的不只是算力。

8.4 补充：PyTorch FSDP（Zhao et al., VLDB 2023）

slides 的 FSDP 部分直接引用这篇。它不提出新的算法（就是 ZeRO-3），价值在于把它做进 PyTorch 核心时的经验：

• 和 PyTorch 核心组件协同设计：Tensor 实现、dispatcher、CUDA 缓存分配器；
• FlatParameter（§5.2）；
• 分片因子 F 和混合分片（§5.5）；
• 单独的通信流、反向预取、前向预取（§5.4）；
• 速率限制器（缓存分配器在 CPU 上做决定带来的问题）；
• 延迟初始化（meta device）；
• 两种 API：FullyShardedDataParallel 包装器和 fully_shard 模块标注（后者保留模块结构和参数名，slides 用的就是它）。

一个值得注意的教训：真实系统的瓶颈常常不在算法里。FSDP 论文用了大量篇幅讨论缓存分配器、CUDA 流、CPU 线程跑得太快——这些在 ZeRO 论文的通信量分析里完全看不到。

8.5 选读：ZeRO-Offload（Ren et al., ATC 2021）

• 问题：只有一块 GPU 的人能不能训几十亿参数的模型？
• 方法：数据流图划分推出唯一最优策略（§7.1），CPU-Adam，一步延迟更新（§7.2）；多卡时和 ZeRO-2 结合，每块卡只卸载自己那份分片的优化器状态。
• 结果：单块 V100 训练 13B（PyTorch 最大 1.4B）；10B 模型 40 TFLOPS；10B 模型在 128 块 GPU 上近似线性扩展；和模型并行结合，一个 DGX-2 上训 70B。
• 和 ZeRO-Infinity 的关系：ZeRO-Infinity 把 FP16 参数也卸载出去（§7.1 枚举里 6M 的那种划分），并用以带宽为中心的分片解决了 ZeRO-Offload “参数只经一条 PCIe” 的问题。

9. 小结 · Summary

1. 显存账：混合精度 Adam 每个参数 16 字节（ZeRO 的记法）或 20 字节（再算 FP32 梯度），其中 12 字节是只在更新时才用到的优化器状态。“放得下 FP16 权重”差得远。
2. ZeRO 的三个阶段：切优化器状态（4Ψ + 12Ψ/N）→ 再切梯度（2Ψ + 14Ψ/N）→ 再切参数（16Ψ/N）。前两个阶段有一份完整 FP16 参数的”地板”，只有第三阶段随卡数线性下降。代码复现了论文表 1、表 2。
3. 通信：前两个阶段 2Ψ，和数据并行一样（AllReduce 本来就是 reduce-scatter + all-gather）；第三阶段 3Ψ（1.5 倍）。slides 的第 1 阶段写法是 3Ψ，因为 AllReduce 后半段的梯度 all-gather 被浪费了。
4. ZeRO 是精确的：代码里五种做法训练 50 步后参数逐位相同；前提是同样的求和顺序。忘了更新后的 all-gather 不会报错，只会让各卡悄悄地不一致。
5. FSDP：FlatParameter（拼接、补零、均分）让一个单元一次通信；单元大小在峰值显存和通信次数之间权衡；单独的通信流 + 反向预取把通信藏进计算；速率限制器防止 CPU 跑太快把显存占满；分片因子 F 在显存和跨节点流量之间连续可调。
6. ZeRO-R：激活分片、固定大小的缓冲区、碎片整理。模型状态切小以后，激活和碎片就成了新瓶颈。
7. 卸载：ZeRO-Offload 的唯一最优划分只把 FP16 参数留在 GPU（可以枚举验证）；ZeRO-Infinity 用 AIT 算清带宽需求，再用以带宽为中心的分片和重叠让 CPU / NVMe 卸载变得足够快，一个节点能训万亿参数。

10. 常见误区 · Common misconceptions

• “ZeRO 是一种近似，会影响收敛。” 不会。ZeRO-1/2/3 和数据并行做完全相同的计算，代码里训练 50 步后参数逐位相同。（ZeRO-Offload 的 DPU 才改变了训练语义，而且是可选的。）
• “ZeRO 一定增加通信。” 第 1、2 阶段的通信量和数据并行完全一样（2Ψ），只有第 3 阶段多 50%。
• “ZeRO-1 和 ZeRO-2 显存差不多。” 差的是完整的梯度缓冲区：N 大时第 1 阶段 4Ψ、第 2 阶段 2Ψ。区别不在算法，而在什么时候释放——第 2 阶段在反向进行中就把不归自己管的梯度释放了。
• “ZeRO-3 / FSDP 每块卡只需要 16Ψ/N 的显存。” 那是稳态。峰值还要加上同时拼完整的单元（以及完整的梯度）。单元切得太大，或者 CPU 跑得太快、没有限速，峰值会远远超过稳态。
• “单元越小越好。” 单元越小峰值越低，但集合通信次数越多，每次都要付 (W − 1) 轮的延迟；消息太小时带宽也用不满。
• “全分片一定比混合分片好，因为最省显存。” 全分片的三次集合通信都跨节点，在慢网络上可能藏不住。只要放得下，混合分片常常更快。
• “卸载到 CPU 一定很慢。” 取决于 AIT 和带宽。优化器状态的更新可以在所有 CPU 上并行，激活检查点的 AIT 很高；参数 / 梯度要高带宽，但用 all-gather 让所有 PCIe 一起工作后，有效带宽随卡数线性增长。

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

练习 1　一个 13B 参数的模型，64 块 80 GB 的卡，混合精度 Adam（K = 12）。分别算数据并行、ZeRO-1、ZeRO-2、ZeRO-3 每块卡的模型状态。哪些放得下？

答案：

• 数据并行 16 × 13 = 208 GB；
• ZeRO-1：4 × 13 + 12 × 13/64 = 52 + 2.44 = 54.4 GB；
• ZeRO-2：2 × 13 + 14 × 13/64 = 26 + 2.84 = 28.8 GB；
• ZeRO-3：16 × 13/64 = 3.25 GB。

只看模型状态，后三个都放得下；但 ZeRO-1 只剩约 25 GB 给激活和临时缓冲区，可能不够。

练习 2　证明：ZeRO-2 每块卡每步发送的字节数和数据并行相同。用环形算法，N 块卡，Ψ 个 FP16 参数。

答案：

• 数据并行：环形 AllReduce = reduce-scatter + all-gather，每块卡各发 (N − 1)/N × 2Ψ 字节，合计 2 × (N − 1)/N × 2Ψ。
• ZeRO-2：梯度 reduce-scatter 发 (N − 1)/N × 2Ψ；更新后 FP16 参数 all-gather 发 (N − 1)/N × 2Ψ，合计相同。
• 代码 02 演示 3 里两者都是 1.500 × 2Ψ（N = 4）。

练习 3　slides 的第 1 阶段写成”反向后 AllReduce 梯度，更新后 all-gather FP16 参数”。这样每块卡每步发多少？哪部分是浪费的？怎样改？

答案：AllReduce 2 × (N − 1)/N × 2Ψ + all-gather (N − 1)/N × 2Ψ = 3 × (N − 1)/N × 2Ψ，即 3Ψ。浪费的是 AllReduce 的后半段：它把求和后的梯度 all-gather 给所有卡，但每块卡只用自己那 1/N。把 AllReduce 换成 reduce-scatter 就回到 2Ψ，显存不变。

练习 4　为什么 ZeRO-3 更新之后不需要 all-gather 参数，而 ZeRO-1、ZeRO-2 需要？

答案：ZeRO-3 的每块卡本来就只存自己那 1/N 的 FP16 参数，下一步前向时才 all-gather——更新后的 all-gather 被推迟、并拆到了每一层的前向里。ZeRO-1、2 每块卡长期存完整的 FP16 参数，自己只更新了 1/N，所以必须马上把其他卡更新的部分拿回来。

练习 5　代码里 ZeRO-3 的峰值是 161,632 字节，稳态（反向结束）是 130,624 字节。各单元的参数数是 [3168, 9312, 9312, 9312, 1552]，N = 4。解释峰值是怎么来的。

答案：峰值出现在反向处理单元 1 时（单元按 4 → 0 的顺序反向）：

• 常驻部分：FP16 参数分片 16,328 + FP32 状态 3 × 32,656 = 114,296；
• 已经 reduce-scatter 完的梯度分片（单元 4、3、2）：(1552 + 9312 + 9312)/4 × 2 = 10,088；
• 单元 1 的完整参数和完整梯度：9312 × 2 × 2 = 37,248。

合计 161,632。所以峰值 ≈ 常驻分片 + 一部分梯度分片 + 一个最大单元的参数和梯度。

练习 6　混合分片：W = 64 卡（8 节点 × 8 卡），F = 8，模型 FP16 大小 M。每块卡每步跨节点发送多少字节？全分片（F = 64）呢？

答案：

• 混合分片：all-gather、reduce-scatter 都在节点内；跨节点只有梯度分片（M/8）在 8 个节点之间的 AllReduce，每卡发 2 × 7/8 × M/8 ≈ 0.22M。
• 全分片：三次集合通信都在 64 卡的大环上，每卡发 3 × 63/64 × M ≈ 2.95M。
• 相差约 13 倍。代码 03 演示 4：5.58 GB vs 75.29 GB（13B 模型，M = 25.5 GB）。

练习 7　用 ZeRO-Infinity 的公式：峰值 70 TFLOPS、序列 1024、每卡 batch 2，参数和梯度要多少带宽才能达到 90% 的效率？优化器状态呢？

答案：

• 参数 / 梯度 AIT = 1024 × 2 = 2048；bw = 0.9 × 70e12 / (0.1 × 2048) ≈ 308 GB/s。
• 优化器状态 AIT = 2048 / 4 = 512；bw ≈ 1.23 TB/s（论文”约 1.5 TB/s”）。

练习 8　（开放）ZeRO-Offload 的 DPU 让第 i + 1 步用第 i − 1 步的梯度更新后的参数。这和第 14 讲会讲的 PipeDream 的”权重暂存”有什么相似之处？为什么论文要先正常训练几十步再打开 DPU？

参考思路：

• 两者都是用过时的参数换取重叠：DPU 让 CPU 更新和 GPU 计算并行，PipeDream 让不同 micro-batch 在不同版本的权重上前进。都会引入”梯度陈旧”（staleness），类似异步 SGD。
• 训练早期梯度变化剧烈，延迟一步的误差大，可能破坏稳定性；几十步之后梯度方向变化平缓，延迟一步的影响很小。论文实测这样做不影响收敛。

12. 延伸阅读 · Further reading

• Y. Zhao et al., PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel（VLDB 2023）——本讲 §5 的主要来源。
• J. Ren et al., ZeRO-Offload: Democratizing Billion-Scale Model Training（ATC 2021）——§7.1–7.2。
• P. Micikevicius et al., Mixed Precision Training（ICLR 2018）——FP32 主参数、损失缩放。
• T. Chen et al., Training Deep Nets with Sublinear Memory Cost（2016）——激活检查点（第 2 讲）。
• V. Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models（MLSys 2023）——序列并行与选择性重计算，激活显存的更精确公式。
• S. Rajbhandari et al., DeepSpeed-MoE（ICML 2022）；以及 DeepSpeed 的 ZeRO++（量化通信、分层分片）——ZeRO 的后续演化。

13. 下一讲预告 · Next lecture

ZeRO 和 FSDP 让数据并行能训任意大的模型——只要卡够多。但它们不切计算：每块卡仍然要算整个模型的前向和反向，而且第 3 阶段要在网络上来回搬参数。第 14 讲《模型并行与流水线并行》换一个方向：把计算切开。

• 张量并行（Megatron-LM）：把一个矩阵乘法切到几块卡上，每层只通信激活，不通信参数；
• 流水线并行（GPipe、1F1B、交错 1F1B、PipeDream）：把模型按层切成几段，像流水线一样传递 micro-batch，代价是”空泡”。

我们会在模拟设备上验证 Megatron 的 MLP 和注意力切分在数值上和单卡相同，数一数每层要几次 AllReduce，再写一个流水线调度模拟器，看 GPipe、1F1B、交错 1F1B 的空泡率和激活显存。