第 18 讲 · 分布式训练中的显存优化：ZeRO · Memory Optimization in Distributed Training: ZeRO

对应 Maps to：CMU 11-868 / 11-968 Large Language Model Systems（Spring 2026，Lei Li），Lecture 18 · Mon 3/23 · Memory Optimization in Distributed Training（slides：llmsys-18-zero.pdf，约 76 页，大部分是逐步动画）。 阅读材料：Rajbhandari et al. (2020) ZeRO: Memory Optimizations Toward Training Trillion Parameter Models（SC）；Ren et al. (2021) ZeRO-Offload；Rajbhandari et al. (2021) ZeRO-Infinity；Wang et al. (2024) ZeRO++（ICLR）；Zhao et al. (2023) PyTorch FSDP（VLDB）。 依据说明：slides 覆盖：数据并行与模型并行的优劣、混合精度 DDP 的显存构成（参数、梯度、优化器状态，20 字节 / 参数的算法）、LLaMA-3 8B 与 GPT-3 175B 的账、临时缓冲与碎片、降低显存的几类方法（ZeRO、激活检查点与压缩、CPU 卸载、省显存的优化器）、ZeRO-1 / 2 / 3 的逐步动画（reduce-scatter、逐分区 reduce、逐段 broadcast）、各阶段的显存公式与通信量、ZeRO-R（分区的激活检查点、固定大小的缓冲区、碎片整理）、ZeRO++、显存与吞吐的基准、与流水线 / 张量并行的组合（3D 并行）、DeepSpeed 代码示例。16 字节 / 参数的算法、图 1 的数字、激活的估计、FSDP 取自论文；HW5 Problem 3 取自作业说明。本讲义的 ZeRO 模拟器、PyTorch 的实测（ZeroRedundancyOptimizer、FSDP2、激活检查点、CPU Adam）、显存与通信的各种表格、HW5 部件的实现与”引擎语义”的模拟，都是本讲义自己设计的，不代表课堂上讲了什么。 配套代码 Code：code/lecture18/（5 个脚本 + zero.py；04 约 12 秒，其余 5 秒以内） 预计阅读时间：2.5 小时（含运行代码）

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

0.1 学习目标 Learning goals

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

1. 算出训练一个模型需要多少显存：模型状态（16 字节 / 参数，以及 slides 的 20 字节算法）、激活、临时缓冲、碎片。
2. 说清楚 ZeRO-1 / 2 / 3 各切开了什么、每块卡还剩多少（Ψ、Nd、K 的公式），以及为什么 ZeRO 不改变训练结果。
3. 推导各阶段的通信量：ZeRO-1 / 2 与数据并行一样是 2Ψ，ZeRO-3 是 3Ψ。
4. 用 PyTorch 的 ZeroRedundancyOptimizer 和 FSDP2 实现 ZeRO-1 / ZeRO-3，知道它们与 DeepSpeed 的区别。
5. 解释激活检查点、分区的激活检查点、CPU 卸载、ZeRO++ 各在省什么、代价是什么。
6. 完成 HW5 Problem 3 的几个部件（DeepSpeed 配置、LoRA、梯度累积的引擎语义、只保存适配器），并估算 LoRA + ZeRO 的显存。

0.2 先修与本讲范围 Prerequisites & scope

0.3 本讲地图 Map

§1  显存都花在哪里：模型状态、激活、临时缓冲、碎片                                  ← 代码 01
§2  ZeRO 的三个阶段：切优化器状态、切梯度、切参数；公式与模拟                        ← 代码 02
§3  通信量；PyTorch 里的 ZeRO：ZeroRedundancyOptimizer 与 FSDP                      ← 代码 02、03
§4  ZeRO 之外：激活检查点、分区的激活检查点、CPU 卸载、ZeRO++                       ← 代码 04
§5  HW5：ZeRO-2 + LoRA                                                            ← 代码 05
§6 小结   §7 常见误区   §8 练习   §9 延伸阅读   §10 下一讲预告

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

1.1 数据并行与模型并行的两难

slides 第 5 页的表：模型并行显存效率好（每块卡只放一部分模型），但计算和通信效率差——Megatron 训练 40B 模型时只跑到峰值的 5%（跨机器的张量并行，第 16 讲 §5.6）；数据并行计算和通信效率好，但显存效率差——每块卡都存一整份模型。ZeRO 的目标：数据并行的效率 + 模型并行的显存。

1.2 模型状态：16 字节还是 20 字节

混合精度 + Adam 训练时，每个参数要存（01 号脚本第 1 部分）：

两种算法只差 FP32 梯度：slides 把它算作常驻的优化器状态，ZeRO 论文认为它可以逐段临时转换。论文把优化器状态的字节数记作 K（Adam 是 12），每块卡的模型状态 = (2 + 2 + K)·Ψ。slides 第 12 页的账：LLaMA-3 8B 的参数 16 GB、梯度 16 GB、优化器状态 128 GB，合计 160 GB；GPT-3 175B 合计 3500 GB。一块 80 GB 的卡连 8B 的模型都训练不了。

01 号脚本第 2 部分在 PyTorch 里真正数了一遍（训练一步之后的参数、梯度、optimizer.state）：

纯 FP32 的 Adam 也是 16 字节（4 + 4 + 4 + 4）——和混合精度的 2 + 2 + 12 巧合地相同。混合精度省的是激活和计算时间，不是模型状态。优化器决定了状态的大小：slides 第 16 页的”省显存的优化器”（Adafactor 只存行列统计量、8 位 Adam 把 m 和 v 量化）就在压这一项。

1.3 激活

slides 第 6 页还列了激活：“每层 d × len × b”。它和批大小、序列长度成正比。01 号脚本第 3 部分算 GPT-2 1.5B（48 层、h = 1600、25 个头；序列 1024、批 32）：

ZeRO 论文说这个配置的激活”约 60 GB”，用激活检查点约 8 GB——那是不含注意力分数的粗略估计；按 Korthikanti 的逐项公式（第 16 讲 §5.6）是 287 GB。这个模型头多（25 个）、h 小，5as/h = 80 远大于 34，注意力分数占了大头。不管哪种估计，激活很容易比模型状态还大，重计算是最直接的办法。

1.4 临时缓冲与碎片

slides 第 13 页：

• 临时缓冲：all-reduce、梯度范数这类操作喜欢把所有梯度拼成一个扁平缓冲区（大张量带宽效率高，第 15 讲的分桶也是这个道理）。1.5B 的模型，一个 FP32 扁平缓冲区就要 6 GB。
• 碎片：长寿命的张量（参数、优化器状态）和短寿命的张量（激活、临时缓冲）交错分配，空闲显存被切成小块，总量够却分配不出一块大的。极端情况下浪费 30%（第 11 讲的缓存分配器）。

2. ZeRO 的三个阶段 · The three stages

2.1 冗余在哪里

数据并行时，K 块卡各存一整份模型状态。可是每块卡更新参数时真的需要全部的优化器状态吗？Adam 是逐元素的：更新第 i 个参数只需要第 i 个参数的 m、v、主参数和平均梯度。如果把参数分成 Nd 段，让第 r 块卡只负责第 r 段，它就只需要这一段的优化器状态——其余的 (Nd − 1)/Nd 都是冗余。

ZeRO（Zero Redundancy Optimizer，Rajbhandari et al. 2020，slides 第 17 页）依次去掉三种冗余：

（slides 第 64 页用 M 表示优化器状态字节数，写成 4N + M·N、4N + M·N/K、2N + (2 + M)·N/K、(4 + M)·N/K，是同一组公式。）

02 号脚本第 1 部分代入 ZeRO 论文图 1 的例子（Ψ = 7.5B、Nd = 64、K = 12）：120 GB → 31.4 GB → 16.6 GB → 1.9 GB，与论文一致。slides 说”7B 的模型 120 GB → 30 GB（4 块卡）“，对应的是 ZeRO-3：7.5B、4 块卡时 30.0 GB。

2.2 ZeRO-1：切优化器状态

slides 第 18–41 页的动画，一步一步：

1. 每块卡用完整的 FP16 参数做前向，得到损失；
2. 反向，得到完整的 FP16 本地梯度；
3. reduce-scatter（ncclReduceScatter）：每块卡只拿到自己负责的那一段的平均梯度；
4. 每块卡只更新自己那一段：FP32 的方差、动量、主参数，再拷成 FP16 参数；
5. all-gather（ncclAllGather）：把各段更新好的 FP16 参数拼回每块卡，开始下一步。

注意第 3 + 5 步：数据并行的 all-reduce = reduce-scatter + all-gather（第 14 讲）。ZeRO-1 只是把这两半拆开，在中间插入”只更新自己那一段”。优化器状态（12Ψ，大头）切成 Nd 份，通信一点没多。

2.3 ZeRO-2：切梯度

slides 第 42–50 页：每块卡仍然算出所有参数的梯度，但只存自己负责的那一段。反向是从后往前逐层进行的，某一段（例如 GPU 3 负责的 M3）的梯度一算完，GPU 0、1、2 就把它 reduce（ncclReduce）给 GPU 3，然后立刻删掉；GPU 3 留下总和。接着是 M2、M1、M0。

每块卡只需要一个临时缓冲，存”正在算、还没发出去”的那一段梯度。逐段的 reduce 合起来就是一次 reduce-scatter——通信量不变，梯度又省了 Nd 倍。

2.4 ZeRO-3：切参数

slides 第 51–63 页：参数也切成 Nd 段，平时每块卡只存自己那一段。

• 前向：算第一段的层之前，GPU 0 把这一段参数 broadcast 给所有卡；所有卡算完这几层，GPU 1、2、3 删掉这一段参数；接着 GPU 1 broadcast 第二段……
• 反向：从最后一段开始，负责的卡再 broadcast 一次参数（前向之后已经删了）；所有卡算出这一段的本地梯度，reduce 给负责的卡（同 ZeRO-2），然后删掉参数和梯度。

每块卡上的参数只有自己那 1/Nd，加上”正在用的那一段”。模型状态随卡数线性下降——这是 ZeRO-1、2 做不到的（它们有一个不随卡数下降的底：4Ψ 和 2Ψ）。

2.5 模拟：结果一样，存的不一样

zero.py 在一个进程里模拟 4 块卡训练一个 14 层的 MLP（Ψ = 52552，Adam，10 步），把每次集合通信和每块卡存的东西都记下来。02 号脚本第 2 部分：

• 训练出的参数逐位相同。ZeRO 不改变任何数学：每个参数还是用所有卡的平均梯度、同样的 Adam 公式更新，只是”谁来存、谁来算”变了。（模拟里集合通信按固定的 rank 顺序求和；真实的环形算法求和顺序不同，会差在最后几位——§3.3 的 FSDP 就是这样。）
• 数据并行、ZeRO-1 与公式完全一致。ZeRO-2、3 多出的一点是临时缓冲：一层的梯度要先完整地算出来才能发走（最大的一层 4160 个元素）；ZeRO-3 还要临时拼出一整层的参数。所以 DeepSpeed 按固定大小的”桶”做通信、预取下一层——零头的大小受控。

2.6 能训练多大的模型

只看模型状态，32 GB 的 V100 上（02 号脚本第 4 部分）：

数据并行加多少卡都是 2B；ZeRO-1、2 的上限趋近 8B、16B；ZeRO-3 随卡数线性增长，1024 块卡时超过 2T——slides 第 71 页”用 1024 块 V100 训练万亿参数的模型”（还要加上激活和缓冲的优化，§4）。slides 第 72 页的实测：ZeRO 训练了 17B 的 Turing-NLG（2020 年 1 月最大的模型），在 400 块卡上训练 100B 的模型，吞吐约是基线的 10 倍、达到理论峰值的 30%。

3. 通信量与 PyTorch 里的 ZeRO · Communication & PyTorch

3.1 通信量

slides 第 65 页：ZeRO-1、2 不增加通信；ZeRO-3 是 1.5 倍。02 号脚本第 3 部分数出模拟里每一步的集合通信（每块卡发送，以 Ψ 为单位，K = 4 时乘 3/4）：

去掉 (K − 1)/K 就是：数据并行、ZeRO-1、ZeRO-2 都是 2Ψ，ZeRO-3 是 3Ψ——多出来的一份是反向时再拼一次参数（前向之后已经释放了）。slides 把这些 all-gather 画成”负责的卡依次 broadcast 自己那一段”，通信量相同。注意次数：ZeRO-2、3 把一次大通信拆成很多次小通信，延迟项（第 14 讲的 α）变多——又是第 15 讲的分桶问题。

为什么 1.5 倍的通信可以接受？因为数据并行的通信本来就能和计算重叠（第 15 讲），而 ZeRO-3 的参数 all-gather 可以预取（算第 l 层时就开始拼第 l + 1 层），在带宽足够时也能大部分藏住。slides 第 73 页的总结：缺点是”某些阶段引入额外的通信，取决于基础设施（PCIe / NVLink）”。

3.2 PyTorch：ZeroRedundancyOptimizer（ZeRO-1）

torch.distributed.optim.ZeroRedundancyOptimizer 包住一个普通的优化器，配合 DDP 使用：每个 rank 只为一部分参数保存状态、只更新这部分，更新后把参数广播给其他 rank。03 号脚本第 1 部分在本机起 4 个进程（gloo），模型 281864 个参数：

四个 rank 的状态加起来正好是一份；训练 5 步后与 DDP + Adam 逐位相同（梯度仍由 DDP 的 all-reduce 平均，没有任何求和顺序变化）。但分得不均：PyTorch 以”整个参数张量”为单位分配（按大小贪心地分给最轻的 rank），rank 0 负责的是平均的 1.16 倍。DeepSpeed 把所有参数拼成一个扁平向量再均匀切开（zero.py 的做法），更均匀，代价是一个张量可能跨两个 rank。

3.3 PyTorch：FSDP（ZeRO-3）

FSDP（Fully Sharded Data Parallel，Zhao et al. 2023）是 PyTorch 版的 ZeRO-3。新的 FSDP2 接口 fully_shard(module) 把模块的参数变成 DTensor，沿第 0 维切成 world_size 份。03 号脚本第 2 部分对每个 Linear 调用 fully_shard：形状 (256, 256) 的权重在每个 rank 上是 (64, 256)；每个 rank 的参数和 Adam 状态都是 1/4（70466 ≈ 281864 / 4）。前向 / 反向进入一个 Linear 时 all-gather 出完整的权重，用完立刻释放。

训练 5 步后与 DDP + Adam 差 5.8e-6（float32）；用 --float64 重跑，差 5.5e-15——是 reduce-scatter 与 all-reduce 的求和顺序不同带来的舍入，不是算法的差别。实践中 FSDP 的”单元”通常是一个 Transformer 层（每层一次 all-gather），并且预取下一层，让通信和计算重叠。

⏱ 本机一步：DDP 约 4 ms、ZeRO-1 约 4–5 ms、FSDP2 约 16–17 ms。模型很小，耗时主要是 gloo 每次集合通信的固定开销，FSDP2 调用次数最多——CPU + gloo 上的时间不代表 GPU + NCCL。

3.4 ZeRO 与 3D 并行

slides 第 74–75 页：ZeRO 也可以和流水线、张量并行组合。常见的搭配是ZeRO-1 + 3D 并行：张量、流水线并行切模型，数据并行维度上再用 ZeRO-1 切优化器状态（Megatron 的”分布式优化器”就是这个，第 16 讲 §6.2 的显存瓶颈正是它解决的）。ZeRO-2、3 和流水线并行一起用比较别扭：流水线的每个微批都要用梯度、参数，频繁的 reduce / gather 会和微批调度打架。

4. ZeRO 之外 · Beyond ZeRO

slides 第 16 页列了几类方法：ZeRO；减少激活（检查点、压缩）；CPU 卸载；省显存的优化器。ZeRO 论文本身还有一部分叫 ZeRO-R，处理”剩余的”显存：激活、缓冲、碎片（slides 第 67–69 页）。

4.1 激活检查点（PyTorch 实测）

第 16 讲 §4.3 用手写的反向讲过原理。PyTorch 里一行就够：checkpoint(block, h, use_reentrant=False)。04 号脚本第 1 部分用 saved_tensors_hooks 数”为反向存下的激活”（12 个块、输入 64 × 128 × 256）：

不用检查点时每个块要存 LayerNorm 的输入和统计量、两个 Linear 的输入、GELU 的输入；每个块一个检查点只存块的输入，反向到这个块时重跑一遍前向。代价约是多一次前向（理论上前向占前向 + 反向的 1/3）。检查点越疏存得越少，但重算时的临时工作区越大（这里没有计入）。

4.2 分区的激活检查点（P_a）

张量并行时，每层的检查点（s × b × h 的输入）在 t 块卡上是完全相同的副本。ZeRO-R 把它切成 t 份，每块卡存 1/t，反向重算这一层之前 all-gather 回来。04 号脚本第 2 部分（100B 的模型：125 层、h = 8192、序列 1024、批 32）：检查点每块卡 62.5 GiB，t = 16 时切开只剩 3.9 GiB，代价是反向时每层 all-gather 512 MiB。这和第 16 讲的序列并行是同一个想法：张量并行组里被复制的激活都可以切开。还可以进一步卸载到 CPU。

4.3 固定大小的缓冲区与碎片整理

• 固定大小的缓冲区（slides 第 68 页）：不再把所有梯度拼成一个随模型增长的扁平缓冲区，而是用固定大小的缓冲区一段一段地做——和 DDP 的分桶一样，模型再大缓冲区也不变。
• 碎片整理（slides 第 69 页）：给长寿命的张量（参数、优化器状态、激活检查点）预先分配一块连续的显存，短寿命的张量（重算出来的激活、临时梯度）在别处分配和释放，互不交错。还可以更进一步，像 LightSeq（第 11 讲）那样按生命周期复用显存。

4.4 CPU 卸载：ZeRO-Offload

把优化器状态和参数更新挪到 CPU 上：GPU 上只留 FP16 的参数和梯度（4Ψ），CPU 上存 FP32 的主参数、m、v（12Ψ），并在 CPU 上做 Adam。每一步：梯度从 GPU 传到 CPU（2Ψ），CPU 更新，FP16 参数传回 GPU（2Ψ）。04 号脚本第 3 部分：

一块卡就能放下 13B 的模型状态。代价是 PCIe 传输和 CPU 上的 Adam（本机约每秒 15 亿个参数）——slides 说”可能占 50% 的时间”。两者都正比于 Ψ，而 GPU 的计算正比于 Ψ × token，所以”每步要多少 token 的计算才能藏住它们”与模型大小无关：按本机的 Adam 速度、A100 40% MFU，每块卡每步约 17K 个 token。大批量的预训练做得到，小批量的微调就藏不住。ZeRO-Offload 的对策是多线程、向量化的 CPU-Adam，以及”延迟一步的参数更新”（CPU 更新和 GPU 的下一步计算重叠）。ZeRO-Infinity 更进一步，把状态卸载到 NVMe。

4.5 ZeRO++：压缩通信

ZeRO-3 的三次通信各传一份 FP16 的模型（3M）。ZeRO++（Wang et al. 2024，slides 第 70 页）三管齐下（04 号脚本第 4 部分）：

跨机器的通信降到 1/4。量化为什么要分块？一个缩放系数管整个张量时，少数大值把它撑大，其余的值都被量化得很粗：我们量化 100 万个权重（夹杂 64 个 20 倍大的值），INT8 整个张量一个 scale 误差 11.25%、每 256 个一个 scale 只有 0.89%；INT4 是 98.69% vs 14.53%。这正是下一讲的主题。

5. HW5：ZeRO-2 + LoRA

HW5 Problem 3 用 DeepSpeed 的 ZeRO-2 引擎做 LoRA 微调。05 号脚本实现了除 DeepSpeed 引擎本身以外的部件，并用一个小引擎模拟它管梯度累积的语义。

5.1 配置（3.1）

{"train_micro_batch_size_per_gpu": micro_batch,
 "gradient_accumulation_steps": accumulation,
 "train_batch_size": world_size * micro_batch * accumulation,
 "zero_optimization": {"stage": 2},
 "zero_allow_untested_optimizer": True,
 "gradient_clipping": 0}               # fp16 时再加 "fp16": {"enabled": True}

train_batch_size 是”一次参数更新看到的样本数”，DeepSpeed 会检查它等于三者之积。ZeRO-2 在配置里只是一行，切分的工作全在引擎里——模型代码不用改。输入小于 1、或者 precision 不是 fp32 / fp16（包括 bf16）要抛 ValueError。

5.2 LoRA（3.2）

LoRA 冻结原来的权重 W₀，只训练一个低秩的增量：h = W₀x + b + (α/r)·B·A·x，A 是 r × in，B 是 out × r 且初始化为 0——刚套上时模型的输出完全不变（05 号脚本第 2 部分：差 0）。apply_lora 的要点是先检查完所有目标，再修改：[‘0’, ‘1’] 里 ‘1’ 是 ReLU，如果边检查边改，‘0’ 已经被换掉、所有参数也已经冻结了，模型就停在”改了一半”的状态。LoRA 的原理留到第 23 讲。

5.3 梯度累积的引擎语义（3.3）

DeepSpeed 的 engine 自己管累积：engine.backward(loss) 把 loss 除以累积步数再反向，engine.step() 每 accumulation 次才真正更新一次并清零梯度。05 号脚本第 3 部分用一个模拟这套语义的小引擎（2 个微批累积一步，3 步），与”整批一步”比较：

两个错误都不会报错：前者把第一个微批的梯度清掉了（只用了一半数据），后者把梯度缩小了两次（相当于学习率减半）。

5.4 只保存适配器（3.4）与显存账

微调的成果只是几个小矩阵：只保存 lora_A / lora_B；加载时先检查所有键和形状，再用 copy_ 原地复制（被拒绝的文件不会把模型改掉一半，优化器手里的参数对象也还是同一个）。

LLaMA-2 7B 的显存账（05 号脚本第 5 部分，4 块卡，只算模型状态）：

LoRA 只训练 0.06% 的参数，梯度和 Adam 状态几乎为零，剩下的是冻结的 FP16 基座（13.5 GB）。所以 LoRA + ZeRO-2 的瓶颈是”每块卡一份基座”，要把基座也切开得用 ZeRO-3。HW5 的模型小，ZeRO-2 + LoRA 足够。

6. 小结 · Summary

1. 训练显存 = 模型状态（混合精度 Adam 16 字节 / 参数；slides 算法 20 字节）+ 激活（随批和序列增长）+ 临时缓冲 + 碎片。优化器状态（12Ψ）是模型状态的大头。
2. ZeRO 去掉数据并行的冗余：ZeRO-1 切优化器状态（4Ψ + KΨ/Nd），ZeRO-2 再切梯度（2Ψ + (2 + K)Ψ/Nd），ZeRO-3 再切参数（(4 + K)Ψ/Nd，随卡数线性下降）。7.5B、64 卡：120 → 31.4 → 16.6 → 1.9 GB。
3. ZeRO 不改变数学：模拟中与数据并行逐位相同。通信量：数据并行、ZeRO-1、ZeRO-2 都是 2Ψ（all-reduce 拆成 reduce-scatter + all-gather），ZeRO-3 是 3Ψ。
4. PyTorch：ZeroRedundancyOptimizer = ZeRO-1（按整个张量分配，不均）；FSDP / FSDP2 = ZeRO-3（每个参数沿第 0 维切开，用到时 all-gather）。
5. ZeRO 之外：激活检查点（省 9–60 倍激活，约慢 30%）、分区的激活检查点（再除以 t）、固定大小的缓冲区与碎片整理、CPU 卸载（GPU 只留 4Ψ，代价是 PCIe 与 CPU Adam）、ZeRO++（量化 + 机器内副本，跨机器通信降到 1/4）。
6. HW5：ZeRO-2 只是一行配置；LoRA 先检查后修改、B 初始化为 0；不要在累积窗口里 zero_grad、不要自己缩放损失；只保存适配器。

7. 常见误区 · Common pitfalls

1. “混合精度能省模型状态。” 混合精度 Adam 和纯 FP32 Adam 一样是 16 字节 / 参数（§1.2）。
2. “ZeRO 是一种近似，会影响收敛。” ZeRO 不改变数学，模拟中逐位相同（§2.5）。
3. “ZeRO-1、2 增加了通信。” 它们只是把 all-reduce 拆成两半，通信量相同；只有 ZeRO-3 是 1.5 倍（§3.1）。
4. “ZeRO-2 加更多卡就能训练任意大的模型。” 每块卡仍有完整的 FP16 参数（2Ψ），上限约 16B（32 GB 卡）；随卡数线性下降的只有 ZeRO-3（§2.6）。
5. “ZeRO 解决了所有显存问题。” 它只管模型状态；激活要靠检查点 / 序列并行 / P_a（§1.3、§4）。
6. “ZeroRedundancyOptimizer 和 DeepSpeed 的 ZeRO-1 完全一样。” PyTorch 按整个张量分配，可能不均（§3.2）。
7. “FSDP 的结果和 DDP 应该逐位相同。” reduce-scatter 的求和顺序不同，只差舍入（§3.3）。
8. “CPU 卸载是免费的。” 每步 4Ψ 字节的 PCIe 传输 + CPU 上的 Adam，小批量时藏不住（§4.4）。
9. “梯度累积时每个微批都要 zero_grad。” 引擎在更新之后自己清零；中途清零会丢掉之前的微批（§5.3）。
10. “LoRA + ZeRO-2 能把 70B 塞进一块卡。” 冻结的基座（2Ψ）每块卡一份；要切开基座得用 ZeRO-3 或量化基座（QLoRA，第 23 讲）。

8. 练习 · Exercises

1. 一个 13B 的模型用混合精度 Adam 训练。按 ZeRO 论文的算法，模型状态多少 GB？按 slides 的算法呢？
2. Ψ = 13B、Nd = 16、K = 12：分别算数据并行、ZeRO-1、2、3 每块卡的模型状态。
3. 为什么 ZeRO-1 不增加通信？用”all-reduce = reduce-scatter + all-gather”解释。
4. ZeRO-3 为什么要在反向时再 all-gather 一次参数？如果前向之后不释放，会怎样？
5. 80 GB 的卡、只看模型状态：ZeRO-2 无论多少块卡最多能训练多大的模型？ZeRO-3 在 256 块卡上呢？
6. 一个模型有 3 个参数张量，大小 100、60、40。ZeroRedundancyOptimizer 在 2 个 rank 上会怎样分？各 rank 的 Adam 状态多少个元素？DeepSpeed 的扁平分区呢？
7. 激活检查点”每 k 个块一个”时，存下的激活和重算的临时工作区各怎样随 k 变化？
8. 张量并行 t = 8、每层检查点 512 MiB、80 层：不分区和分区（P_a）时每块卡各存多少？
9. 一个 7B 的模型用 ZeRO-Offload，PCIe 25 GB/s。每步 PCIe 传输多少字节、多少秒？如果每步每卡 8K 个 token、A100 40% MFU，计算多少秒？藏得住吗？
10. ZeRO++ 的 hpZ 用什么换来了”反向不跨机器 all-gather”？
11. HW5：world_size = 2、micro_batch = 2、accumulation = 2，build_config 的 train_batch_size 是多少？如果你在 train_microbatch 里写了 loss = loss / 2，会发生什么？
12. LLaMA-2 13B（40 层、d = 5120、FFN 13824、词表 32000），LoRA r = 16 加在 q、k、v、o 四个投影上：可训练参数多少？4 块卡 ZeRO-2 时每块卡的模型状态约多少？

参考答案

1. ZeRO 论文：16 × 13B = 208 GB；slides：20 × 13B = 260 GB。
2. 数据并行 16 × 13 = 208 GB；ZeRO-1：4 × 13 + 12 × 13/16 = 52 + 9.75 = 61.75 GB；ZeRO-2：2 × 13 + 14 × 13/16 = 26 + 11.375 = 37.4 GB；ZeRO-3：16 × 13/16 = 13 GB。
3. 数据并行的 all-reduce 本来就由 reduce-scatter（每块卡得到一段的总和）和 all-gather（把各段发给所有卡）组成。ZeRO-1 在两者之间插入”只用自己那段的平均梯度更新自己那段参数”，再 all-gather 更新后的参数——两次通信的量和原来的 all-reduce 完全一样。
4. 反向第 l 层要用它的参数（算对输入的梯度：g·Wᵀ）。前向之后已经释放了，所以要再拼一次；如果不释放，前向结束时每块卡就存着完整的参数，失去了 ZeRO-3 的意义。
5. ZeRO-2 上限 80 GB / 2 字节 = 40B（再加优化器状态的 1/Nd 部分，卡越多越接近）。ZeRO-3：80 GB × 256 / 16 = 1.28T。
6. 贪心：100 给 rank 0，60 给 rank 1，40 给当前最轻的 rank 1 → rank 0：100 个参数（Adam 状态 200），rank 1：100（200）——这个例子正好均衡；若大小是 100、30、20，则 rank 0 = 100、rank 1 = 50，状态 200 vs 100。扁平分区：总 200 个，每个 rank 100 个（一个张量可能被切成两半）。
7. 存下的激活 ≈ n/k 个块输入（k 越大越少）；重算时要临时存一段 k 个块的全部中间结果（k 越大越多）。总量 ≈ n/k + k·(每块的激活)，k ≈ √n 时最小（第 16 讲）。
8. 不分区：80 × 512 MiB = 40 GiB；分区：40 / 8 = 5 GiB。
9. 4Ψ = 28 GB → 1.12 s。计算：6 × 7e9 × 8192 / (312e12 × 0.4) ≈ 2.76 s。传输藏得住；CPU Adam 若按本机约 1.5 G 参数 / 秒要约 4.7 s，藏不住——需要更快的 CPU-Adam 或更大的批。
10. 用显存换通信：每台机器内额外存一份参数分片（参数在机器内的 8 块卡上切开，而不是在全部卡上），反向的 all-gather 只在机器内部进行。
11. 2 × 2 × 2 = 8。engine.backward 已经除以累积步数，再除一次，梯度变成应有的 1/2，相当于学习率减半——不会报错，但结果和参照对不上（§5.3）。
12. 每个投影 r × (5120 + 5120) = 16 × 10240 = 163,840；4 个投影 × 40 层 = 26.2M。基座 Ψ ≈ 40 × (4 × 5120² + 3 × 5120 × 13824) + 2 × 32000 × 5120 ≈ 13.0B，FP16 约 26 GB；适配器 16 × 26.2M / 4（ZeRO-2 下优化器状态和梯度切开，适配器的参数 2 字节一份）≈ 0.1–0.4 GB。每块卡约 26 GB。

9. 延伸阅读 · Further reading

• Rajbhandari et al. (2020) ZeRO: Memory Optimizations Toward Training Trillion Parameter Models（SC）：本讲的主要来源（ZeRO-DP 与 ZeRO-R）。
• Ren et al. (2021) ZeRO-Offload: Democratizing Billion-Scale Model Training（ATC）：CPU 卸载、CPU-Adam、延迟一步的更新。
• Rajbhandari et al. (2021) ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning（SC）：NVMe 卸载。
• Wang et al. (2024) ZeRO++: Extremely Efficient Collective Communication for Giant Model Training（ICLR）。
• Zhao et al. (2023) PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel（VLDB）；PyTorch 文档 “FSDP2 / fully_shard”。
• Chen et al. (2016)、Korthikanti et al. (2022)：激活检查点与选择性重计算（第 16 讲）。
• DeepSpeed 教程：“Zero Redundancy Optimizer”、“ZeRO-Offload”；原课的代码示例 llmsys_code_examples/deepspeed_example。

10. 下一讲预告 · Next lecture

ZeRO 用”切开存”省显存；ZeRO++ 已经让我们看到另一条路——用更少的位数存。第 19 讲进入第五单元：模型量化。Llama-70B 推理要 140 GB 显存，量化成 4 位只要 35 GB。可是位数少了，精度从哪里损失？我们会从计算机里的低精度数字讲起（BF16、FP8、INT8、INT4），写出 absmax 和零点量化，比较逐张量、逐通道、逐组量化的误差，然后在一个小 GPT 里植入”离群特征”，看看 LLM.int8() 为什么要把少数几个维度留在 16 位。