起因#
上一篇算完了 7B 模型的训练显存账本。结论很直接：BF16 mixed precision + AdamW 的 Model States 就有 112 GB，加上 30 GB 的 Activations，单卡至少需要 145 GB——一张 80 GB 的 A100 连 Model States 都放不下。
那篇文章的最后提到了两种方向：Activation Checkpointing 和 CPU Offloading。前者把 Activations 从 30 GB 砍到 0.5 GB，后者把 84 GB 的 Optimizer States 搬到 CPU 内存。但当时我跳过了一个问题没细想：
这两种方案都在"省显存"，但省显存的代价完全不同。一个多算一遍 forward，一个多搬几十 GB 数据。什么时候该用哪个？为什么有时候两个都得用？
后来继续看 MIT 6.S191 Lecture 9 关于 Offloading 的部分，才发现这个问题比我想的要复杂——不只是"搬不搬"的问题，还涉及搬多快、能不能和计算重叠、以及搬运本身是否会成为新的瓶颈。
都是省显存，但资源交换不一样#
先把两种方案的核心思路放到一起看。
Checkpointing：Forward 时主动丢弃部分中间 Activation，backward 需要时重新计算。用额外的 FLOPs 换 Activation 显存。
Offloading：把 GPU HBM 中的数据搬到更大但更慢的存储（CPU DRAM 甚至 NVMe），需要时再搬回来。用数据传输换 HBM 容量。
两者的共同点是都不减少数据总量——Checkpointing 没有消灭 Activation 的计算需求，Offloading 没有让 Optimizer States 凭空消失。它们改变的是数据"何时存在"和"存在哪里"。
但代价的性质完全不同。Checkpointing 的代价是计算量，受限于 GPU 算力；Offloading 的代价是数据搬运，受限于 CPU-GPU 之间的传输带宽。这两种资源的瓶颈点和优化方向都不一样。
Checkpointing 只管 Activation 的部分#
关于 Activation Checkpointing 的原理，之前的文章已经详细讨论过：backward 需要 forward 的中间结果，这些中间结果的 lifetime 被 backward 拉得很长。Checkpointing 的做法是只保留少数边界 Activation，backward 到达时从边界重新跑一遍 forward，用完即释放。
用上一篇的 7B 数字：
- Activations（无 Checkpointing）：\(\approx 30.6\) GB
- Activations（full Checkpointing）：\(\approx 0.5\) GB
- 省了约 30 GB，代价是约 33% 额外 FLOPs
但 Model States 的 112 GB 完全不变。
$$ 112\ \text{GB (Model States)} + 0.5\ \text{GB (Activations)} + 2\ \text{GB (overhead)} \approx 114.5\ \text{GB} $$还是超过 80 GB。如果显存的主要压力来自 Model States 而不是 Activations，Checkpointing 帮不上忙。
这件事在 7B 这个量级上特别明显：Model States 112 GB 是 Activations 30 GB 的近四倍。Checkpointing 能把 30 GB 砍到 0.5 GB，但 112 GB 纹丝不动。
84 GB 的 Optimizer 数据在 forward/backward 时干了什么#
写到这里我停下来想了一个问题：这 112 GB 的 Model States 中，哪些数据在 forward 和 backward 阶段真正被使用了？
答案是：只有 BF16 参数（14 GB）和 BF16 梯度（14 GB）。
剩下的 84 GB——FP32 master weights（28 GB）、Adam first moment \(m\)（28 GB）和 second moment \(v\)（28 GB）——在整个 forward 和 backward 过程中完全不参与计算。它们只有到了 optimizer step 时才会被使用。
也就是说，forward 花几百毫秒，backward 又花几百毫秒，这期间 84 GB 的 Optimizer States 只是安静地占着 GPU 显存，什么也没做。
这是一个很自然的问题：既然不用，为什么不搬走？
Offloading 的两种实现路线#
“搬走 Optimizer States"这个思路很直接，但实际有不同的实现方式，数据传输模式也不一样。
变体 A：States 放 CPU，更新在 GPU 做。 Optimizer step 时把 States 分块从 CPU 搬回 GPU，GPU 完成参数更新，再把更新后的 States 搬回 CPU。双向传输。
变体 B：States 放 CPU，更新也在 CPU 做。 Backward 结束后，GPU 把梯度发送给 CPU。CPU 用 FP32 master weights、\(m\)、\(v\) 和收到的梯度完成参数更新，再把更新后的 BF16 参数发回 GPU。这是 ZeRO-Offload 的路线。
两者的区别在于：
- 变体 A 需要把大量 FP32 States 搬回 GPU（84 GB 级别），GPU 完成更新后再搬回去。双向都是大数据量。
- 变体 B 只传输 BF16 梯度（14 GB）到 CPU，再传回 BF16 参数（14 GB）。单次传输的数据量小得多，但 optimizer step 的计算由 CPU 承担。
不能把两种实现混为一谈——它们的传输瓶颈和计算瓶颈位置不同。
后续讨论以变体 B（ZeRO-Offload 路线）为主。在这种实现下，forward/backward 阶段 GPU 上只需要 BF16 参数和 BF16 梯度：
$$ 14\ \text{GB (params)} + 14\ \text{GB (grads)} = 28\ \text{GB} $$GPU 的 Model States 从 112 GB 降到 28 GB。84 GB 的 Optimizer States 转由 CPU DRAM 承担。
搬数据到底需要多久#
GPU 显存是省下来了，但数据搬运不是免费的。
最粗略的估算：
$$ T_{\text{transfer}} \approx \frac{\text{Data Size}}{\text{Effective Bandwidth}} $$CPU 和 GPU 之间通常通过 PCIe 连接。PCIe 4.0 x16 的理论单向带宽约 32 GB/s，但实际的有效带宽会受到协议开销、传输粒度和系统负载的影响，通常低于理论值。
用一个教学假设来建立数量级感觉：
- 传输 20 GB 数据
- 有效带宽 20 GB/s
- 单向传输时间 \(\approx 1\) 秒
如果 optimizer step 需要把 14 GB 梯度发送到 CPU、再接收 14 GB 更新后的参数，累计约 28 GB 的双向传输。在 20 GB/s 有效带宽下约 1.4 秒。
这里有两个容易搞混的地方。
第一，CPU DRAM 的本地带宽和 CPU-GPU 之间的有效传输带宽不是一回事。DDR5 的本地带宽可以到 100 GB/s 以上，但数据要从 CPU DRAM 到 GPU HBM，必须经过 PCIe 通道。PCIe 的带宽通常远低于两端各自的本地带宽。关于 GPU 存储层级的详细讨论可以参考这篇。
第二，传输时间不能直接等同于训练总耗时的增加值。因为传输和计算有可能同时进行。
GPU 一定要干等数据搬完吗#
这是我在看 6.S191 这部分内容时想了比较久的地方。
假设一个训练 step 中，GPU 计算（forward + backward）需要 200 ms，数据传输需要 100 ms。如果串行执行：先算完，再搬数据，总耗时 300 ms。
但 GPU 有独立的 DMA 引擎负责数据传输，和计算单元可以并行工作。如果传输和计算在时间上完全重叠：
$$ T \approx \max(200, 100) = 200\ \text{ms} $$传输的 100 ms 被"隐藏"在了计算的 200 ms 里面。
但换一个场景：计算 200 ms，传输 300 ms。即使完全重叠：
$$ T \approx \max(200, 300) = 300\ \text{ms} $$传输时间比计算时间长，GPU 算完了还要等数据搬完。传输成为了关键路径上的瓶颈。
所以 \(T \approx \max(T_{\text{compute}},\ T_{\text{transfer}})\) 这个公式只在一组前提下成立：
- 计算和传输使用独立的硬件资源（GPU 计算单元 vs DMA 引擎）
- 两者在时间上能够完全重叠
- 没有数据依赖阻止重叠——也就是说，当前的计算不需要等正在传输的数据
这三个条件在真实系统中不一定完全满足。特别是第三个——数据依赖。
为什么要分块#
如果把数据传输当作一个原子操作——先把所有需要的数据一口气搬完，再开始计算——那就是纯粹的串行，没有任何重叠。
要实现重叠，必须分块。把数据拆成若干小块，每搬完一块就可以开始计算这一块对应的结果，同时继续搬下一块。
这和 GPU 做矩阵乘法时的 tiling 是类似的思想：不是把整个矩阵搬到片上再开始算，而是搬一小块、算一小块、搬下一小块——形成流水线。
但分块引入了更多的调度约束。每一块的计算都可能依赖这一块的数据已经到位。如果某一块的传输延迟了，对应的计算就得等。传输时间不能总被完全隐藏。
实际系统中，重叠的程度取决于分块粒度、传输延迟的稳定性、计算密度是否足够"填满"传输等待的间隙，以及 PCIe 带宽是否被其他传输抢占。这些都是工程层面的变量，不是一个简单的 max 公式能概括的。
哪些数据最适合卸载#
不是所有占显存的 tensor 都适合搬到 CPU。我发现区分它们的关键不是"大小”，而是什么时候被访问以及能不能提前搬回来。
Optimizer States#
Optimizer States（master weights + \(m\) + \(v\)）只在 optimizer step 时使用。整个 forward + backward 期间它们不参与计算。
这意味着：从上一次 step 结束到下一次 step 开始，中间有完整的 forward + backward 的时间窗口。可以在这段时间里提前把下一次 step 需要的 States 分块搬回来。
调度窗口大，访问模式集中——Optimizer States 是最适合 offload 的数据。
Activations#
Activations 的情况更复杂。每一层的 Activation 在 forward 时产生，在对应层的 backward 时被使用。
以一个 20 层网络为例。第 1 层的 Activation（Act_L1）在 forward 的开头产生，但要到 backward 的最后（B1）才会被使用。中间隔了 L2 到 L20 的 forward 加上 B20 到 B2 的 backward——这是一个很宽裕的窗口。
如果在 L1 forward 结束后立即把 Act_L1 搬到 CPU，然后在 B2 结束前（也就是 B1 开始之前）提前搬回来，GPU 上这段时间就不需要为 Act_L1 保留显存。
但第 19 层的 Activation（Act_L19）在 F19 产生，B19 就要用。中间只隔了 F20 和 B20 两步——窗口非常短，可能来不及完成一次往返传输。
早期层的 Activation 有更宽裕的卸载窗口，晚期层的窗口可能短到无法完成传输。 这也是为什么 Activation Offloading 在实际系统中通常需要精心调度，而不是简单地"全部搬走"。
Parameters#
参数在 forward 和 backward 中都会被使用。如果把参数 offload 到 CPU，每一层的 forward 和 backward 都需要把对应的参数搬回来。
对于一个 20 层网络，forward 搬 20 次，backward 搬 20 次——总共 40 次传输。而且每次传输和对应层的计算之间有严格的依赖关系：不搬完就不能算。
Parameters 被反复访问，offload 的传输成本最高，通常不是首选。
选型的关键#
决定一个 tensor 是否值得 offload，不只看它有多大，还要看：
- 何时被访问：是集中在一个阶段（Optimizer States），还是分散在整个过程中（Parameters）？
- 调度窗口有多大：从上次使用到下次使用之间，有没有足够时间完成传输？
- 能否提前传输：是否可以异步预取，不阻塞计算？
GPU HBM 之外：CPU DRAM 和 NVMe#
到目前为止讨论的 Offloading 目标都是 CPU DRAM。但 CPU DRAM 的容量也是有限的。
一台典型的训练服务器可能配 256 GB 到 1 TB 的 CPU 内存。7B 模型的 Optimizer States 是 84 GB，放进 256 GB 的 CPU DRAM 没问题。但如果是更大的模型，或者一台机器上跑多个训练任务，CPU DRAM 也可能不够用。
这时候可以继续往下一级存储走：NVMe SSD。
三级存储的粗略对比：
| GPU HBM | CPU DRAM | NVMe SSD | |
|---|---|---|---|
| 典型容量 | 40–80 GB | 256 GB–1 TB | 1–8 TB |
| 带宽（对 GPU 而言） | \(\sim\)2 TB/s（本地） | \(\sim\)25–32 GB/s（PCIe） | \(\sim\)3–7 GB/s（顺序读写） |
| 延迟 | 最低 | 中等 | 较高 |
| 数据路径 | 片上直接访问 | 经 PCIe 到 GPU | 经 NVMe 控制器 → CPU → PCIe → GPU |
几个需要注意的地方：
- 上表的"带宽"指的是从 GPU 视角能有效使用的带宽，不是各存储介质的本地理论带宽。
- NVMe 的数据到 GPU 通常需要经过 CPU 内存中转（除非使用 GPUDirect Storage 等技术直接走 PCIe，这里不展开）。路径越长，延迟越高。
- NVMe 的顺序读写带宽还算不错（PCIe 4.0 x4 的 NVMe 理论约 7 GB/s），但随机读写和小块 I/O 的性能会明显下降。
NVMe Offloading 的容量几乎不受限，但 I/O 代价显著大于 CPU DRAM Offloading。 它更适合那些不需要频繁传输、可以在后台慢慢搬的数据——比如很长时间才访问一次的 Optimizer States 分片。
回到 7B：四种配置的显存账本#
用之前文章的一致数字，整理一下四种配置的 GPU 显存和 CPU 显存。
假设条件：7B 参数，BF16 mixed precision + AdamW（16 bytes/param），\(B = 1\)，\(S = 2048\)，单卡训练。以下数字是教学性的同时驻留估算，不是实测 Peak Memory。
| 方案 | GPU 显存 | CPU 内存 | 额外性能代价 |
|---|---|---|---|
| A. 原始训练 | \(\approx 145\) GB | 0 | 无 |
| B. Activation Checkpointing | \(\approx 114.5\) GB | 0 | \(\sim\)33% 额外 FLOPs |
| C. Optimizer Offloading | \(\approx 60.6\) GB | 84 GB | PCIe 传输开销 |
| D. 两者组合 | \(\approx 30.5\) GB | 84 GB | 两者代价叠加 |
必须说明的是：这些数字忽略了临时缓冲区、重计算时的 workspace、传输 staging memory 等开销。真实的 Peak Memory 会更高。
几个观察：
方案 A 和 B 都放不进 80 GB。 Model States 112 GB 已经超标。Checkpointing 只减了 30 GB Activations，于事无补。
方案 C 能放下，但 Activations 仍然占 30 GB。 如果增大 batch size 或 sequence length，Activations 会线性甚至二次增长（上一篇算过，\(S\) 翻倍时 Activations 涨 3.4 倍），很快又放不下了。
方案 D 既 offload 又 checkpoint，余量最大。 但性能代价也最大：既要多算 33% 的 FLOPs，又要承受 PCIe 传输延迟。
一个值得想的问题#
如果方案 C（只 Offloading）已经放得下（60.6 GB \(\lt\) 80 GB），为什么还要加 Checkpointing？
因为 60.6 GB 中有 30.6 GB 是 Activations。在 \(B = 1\)、\(S = 2048\) 时还好。但训练通常会用更大的 micro-batch：\(B = 4\) 时 Activations 涨到 122 GB，单靠 Offloading 又放不下了。
加上 Checkpointing（方案 D），Activations 降到 0.5 GB 量级，不再是瓶颈，batch size 和 sequence length 有更大的调整空间。
最省显存的方案不一定是吞吐量最优的方案。 方案 D 余量最大，但计算和传输代价都有。如果当前 batch size 下方案 C 就够用，加 Checkpointing 只是增加了 33% 的 FLOPs，得不偿失。
选型不是看"哪个能省更多显存"，而是看"在当前瓶颈下，哪种代价更小"。
训练 OOM 了，应该先做什么#
把前面讨论的内容收成一个工程判断思路。
第一步：搞清楚显存花在哪了。 是 Model States 占大头，还是 Activations 占大头？不同的瓶颈对应不同的优化路线。
第二步：如果 Activations 是主要压力，优先评估 Checkpointing。 代价是约 33% 额外 FLOPs，但不需要引入 CPU-GPU 数据传输，实现也相对简单。
第三步：如果 Model States 是主要压力（特别是 Optimizer States），评估 CPU Offloading。 代价是 PCIe 传输开销。需要评估传输量、传输时间和能否与计算重叠。
第四步：如果容量极度紧张，两者组合使用。 或者考虑更大容量的存储层（NVMe），但要接受更高的 I/O 代价。
第五步：评估性能。 测量 training step time、GPU utilization、PCIe 吞吐和等待时间。单纯让 OOM 消失不是目的——如果训练速度掉了一倍，可能还不如减小 batch size 或者换更大显存的 GPU。
这不是对所有训练系统都适用的硬性规则。不同的模型规模、hardware 配置和训练预算，最优选择可能完全不同。关键是先定位瓶颈，再选策略。
写在最后#
回过头看这篇和前面几篇讨论 GPU 显存的文章，有几件事逐渐变得清楚了。
Checkpointing 和 Offloading 解决的是不同类型的显存压力。 前者减 Activation，后者减 Model States。用错了方向，效果有限。
显存优化的本质是资源交换。 Checkpointing 用计算换显存，Offloading 用数据搬运换显存。没有免费的优化——每种方案都有性能代价，区别在于代价落在哪种资源上。
GPU 显存有限，不代表所有训练数据都必须永久驻留 GPU。 Optimizer States 在 forward/backward 阶段完全不参与计算，Activations 只在特定阶段被使用。真正需要优化的是数据在何时、何处保存，以及如何让数据在被使用时及时到位，同时尽量不增加关键路径上的等待时间。
到这里，讨论还局限在单张 GPU 上的优化。但 7B 模型即使用了 Checkpointing + Offloading，训练速度也会受到 PCIe 传输和额外 FLOPs 的影响。当这些单卡优化的代价已经不可接受时，下一步自然是考虑让多张 GPU 分担训练状态——不过那是另一个话题了。
参考资料#
- 7B 模型只有 14 GB 权重，一张 80 GB GPU 为什么仍可能 OOM？（本文前篇）
- 从推理到训练：模型调度中的显存约束
- 从参数量到 Peak Memory：模型显存到底应该怎么算？
- GPU 的数据存在哪？——从 HBM 到片上存储的存储层级
- 从 CS336 到 ML Infra：理解 Arithmetic Intensity
- MIT 6.S191: Introduction to Deep Learning — Lecture 9: Secrets to Massively Parallel Training. https://introtodeeplearning.com/
- Rajbhandari et al. ZeRO-Offload: Democratizing Billion-Scale Model Training. USENIX ATC 2021.
- Rajbhandari et al. ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning. SC 2021.
- Chen et al. Training Deep Nets with Sublinear Memory Cost. 2016.