起因#
之前分别写了推理和训练的显存心智模型和混合精度训练中的显存组成，反复出现一个粗估：BF16 mixed precision + AdamW 大约 16 bytes/parameter。但这个数字每次都只是"1B × 16 = 16 GB"的单位换算，没有真正对着一个具体的 Transformer 架构把账从头到尾算一遍。
最近看 MIT 6.S191 Lecture 9 讲 Training Memory Model，提到一个问题让我停下来想了想：7B 模型 BF16 权重只有 14 GB，一张 80 GB 的 A100 为什么还可能 OOM？
我第一反应是应该放得下。14 GB 到 80 GB 中间有 66 GB 的余裕，怎么也够了。
但把账算完以后发现，连 Model States 都已经超过 80 GB 了，Activations 还没开始算。
先给 7B 模型一个具体 spec#
不讲概念，直接给一个 7B 级别 Transformer 的架构参数（参考 LLaMA-2 7B 量级）：
| 超参 | 值 |
|---|---|
| 参数量 | \(\approx 7\text{B}\) |
| Hidden dim \(d\) | 4096 |
| Layers \(L\) | 32 |
| Attention heads \(H\) | 32 |
| Head dim \(d_h = d / H\) | 128 |
| FFN intermediate \(d_{\text{ff}}\) | \(4d = 16384\) |
| Sequence length \(S\) | 2048 |
| Micro-batch size \(B\) | 1 |
BF16 下，权重占用：
$$ 7 \times 10^9 \times 2\ \text{bytes} = 14\ \text{GB} $$14 GB，这就是标题里的数字。一张 80 GB 的 A100 看起来绰绰有余。
Model States：从 14 GB 到 112 GB#
训练不只需要权重。BF16 mixed precision + AdamW 的配置下，每个参数对应五份数据。之前的文章已经详细解释了每一项为什么存在，这里只列数字。
| 组件 | 精度 | bytes/param | 7B 占用 |
|---|---|---|---|
| Parameters | BF16 | 2 | 14 GB |
| Gradients | BF16 | 2 | 14 GB |
| Master Weights | FP32 | 4 | 28 GB |
| Adam first moment \(m\) | FP32 | 4 | 28 GB |
| Adam second moment \(v\) | FP32 | 4 | 28 GB |
| Model States 合计 | 16 | 112 GB |
几句话回顾一下这些项的来源（不展开推导，展开见前文）：
Gradients：backward 为每个参数计算 \(\partial L / \partial \theta\)，需要和参数同样大小的存储。（Training Loop 那篇详细讲了为什么 backward 不更新模型、gradient 存在哪。）
Master Weights：BF16 只有 7 位 fraction，参数更新步长如果小于相邻可表示数的间距，更新会被舍入吞掉。所以需要一份 FP32 副本来累积微小更新。（Mixed Precision 那篇用 \(w = 1.25\)，\(\Delta w = -0.001\) 做了完整的数值例子。）
Adam \(m\) 和 \(v\)：AdamW 为每个参数维护一阶矩和二阶矩的滑动平均，各需要一份 FP32 tensor。（Training Loop 那篇解释了 Adam 的数学和它对显存的影响。）
仅 Model States 就需要 112 GB，已经超过了一张 80 GB A100 的全部显存。 还没开始算 Activations。
需要强调：16 bytes/parameter 是特定配置（BF16 参数/梯度 + FP32 master weights + FP32 AdamW）下的估算。换 optimizer（SGD 没有 \(m\)、\(v\)），或者换 precision scheme（FP32 参数本身就是 master），或者加上 ZeRO sharding，单卡 footprint 都会变。这里先不展开这些优化手段。
Activation Memory：一个 Transformer Block 到底保存了什么#
之前的文章已经讲过一个定性结论：backward 改变了 activation 的 lifetime，很多 forward 中间结果必须一直活到对应的 backward pass。但"activation 很大"到底是多大？我之前一直没有算过一个 Transformer Block 的 forward 到底产生多少 bytes 的 saved activation。
6.S191 的 lecture 给出的思路是：逐项拆解 Transformer Block 中每一步计算需要为 backward 保留的 tensor，然后把它们加起来。
逐项拆解#
一个标准 Transformer Block 的 forward 中，autograd 需要保存以下中间结果：
Self-Attention 部分：
| 保存内容 | Shape | 精度 | 每层占用 |
|---|---|---|---|
| LayerNorm 输入 | \(B \times S \times d\) | BF16 | 16 MB |
| Q projection 输出 | \(B \times S \times d\) | BF16 | 16 MB |
| K projection 输出 | \(B \times S \times d\) | BF16 | 16 MB |
| V projection 输出 | \(B \times S \times d\) | BF16 | 16 MB |
| Attention Score \(QK^T\!/\!\sqrt{d_h}\) | \(B \times H \times S \times S\) | BF16 | 256 MB |
| Softmax 输出 | \(B \times H \times S \times S\) | BF16 | 256 MB |
| Attention Dropout mask | \(B \times H \times S \times S\) | bool | 128 MB |
| Attention 输出 \(\text{Softmax} \times V\) | \(B \times S \times d\) | BF16 | 16 MB |
| Output Projection 输出 | \(B \times S \times d\) | BF16 | 16 MB |
| Residual Dropout mask | \(B \times S \times d\) | bool | 8 MB |
FFN 部分：
| 保存内容 | Shape | 精度 | 每层占用 |
|---|---|---|---|
| LayerNorm 输入 | \(B \times S \times d\) | BF16 | 16 MB |
| Up Projection 输出 | \(B \times S \times 4d\) | BF16 | 64 MB |
| GeLU 输出 | \(B \times S \times 4d\) | BF16 | 64 MB |
| Down Projection 输出 | \(B \times S \times d\) | BF16 | 16 MB |
| Residual Dropout mask | \(B \times S \times d\) | bool | 8 MB |
最右列的数字代入了 \(B = 1\)，\(S = 2048\)，\(d = 4096\)，\(H = 32\)。
看这张表有一个地方很不直观：Attention Score 矩阵一项就占 256 MB，是单个 \(B \times S \times d\) tensor 的 16 倍。 加上 Softmax 输出和 Dropout mask，仅 attention score 相关的三项就占了 640 MB，超过单层总量的 70%。
公式#
把表中的项合并，可以得到一个简洁的公式。
所有 \(B \times S \times d\) 形状的 tensor（BF16 按 2 bytes，bool 按 1 byte 计）合并后系数是 34：
$$ 8 \times 2 + 2 \times 1 + 2 \times 8 = 34 $$其中 8 项 BF16 的 \(B \times S \times d\) tensor 贡献 16，2 项 dropout mask 贡献 2，2 项 \(B \times S \times 4d\) 的 FFN tensor 贡献 16（\(4d\) 展开后系数 ×4，每项 2 bytes，共 \(2 \times 2 \times 4 = 16\)）。
所有 \(B \times H \times S \times S\) 形状的 tensor 合并后系数是 5（两项 BF16 各 2，一项 bool 各 1）。
于是单层 activation 显存：
$$ M_{\text{act/layer}} = BSd \times \left(34 + \frac{5HS}{d}\right) \ \text{bytes} $$这个公式来自 Korthikanti et al. 2022（Megatron-LM 团队），假设 \(d_\text{ff} = 4d\)、标准 Multi-Head Attention、BF16 存储。
代入 7B 的参数#
$$ \frac{5HS}{d} = \frac{5 \times 32 \times 2048}{4096} = 80 $$$$ M_{\text{act/layer}} = 1 \times 2048 \times 4096 \times (34 + 80) = 1 \times 2048 \times 4096 \times 114 \approx 0.96\ \text{GB} $$32 层总计：
$$ M_{\text{act}} = 32 \times 0.96 \approx 30.6\ \text{GB} $$仅 Activations 就有约 30 GB。 加上 Model States 的 112 GB，单卡训练 7B 模型至少需要 142 GB 以上的显存——接近 80 GB A100 的两倍。
Batch size 和 Sequence length 的影响#
公式里有两项都和 workload 相关：\(BSd\) 项随 batch size 线性增长，\(BHS^2\) 项还额外随 sequence length 二次增长。
代入不同配置看一下：
| 配置 | \(B\) | \(S\) | 单层 Activation | 32 层合计 |
|---|---|---|---|---|
| Baseline | 1 | 2048 | 0.96 GB | 30.6 GB |
| Larger batch | 4 | 2048 | 3.82 GB | 122 GB |
| Longer sequence | 1 | 4096 | 3.25 GB | 104 GB |
| Both | 4 | 4096 | 13.0 GB | 417 GB |
两个观察：
Batch size 翻倍，activation 翻倍。 线性关系，没有意外。
Sequence length 翻倍，activation 涨了 3.4 倍——因为 attention score 矩阵是 \(S \times S\) 的，它的显存随 \(S\) 二次增长。在 \(S = 2048\) 时，这个二次项已经贡献了单层 activation 的 70%。sequence length 越长，这个比例越大。这也是为什么长序列训练对显存的压力远大于简单地增大 batch size。
完整账本#
把 Model States、Activations 和运行时开销合到一起：
| 组件 | 占用 |
|---|---|
| BF16 Parameters | 14 GB |
| BF16 Gradients | 14 GB |
| FP32 Master Weights | 28 GB |
| FP32 Adam \(m + v\) | 56 GB |
| Activations（\(B\!=\!1, S\!=\!2048\)） | 30.6 GB |
| CUDA context + workspace | ~1–2 GB |
| 合计 | ~144 GB |
回到标题的问题：14 GB 权重，80 GB 显存，中间的 66 GB 去哪了？
- 梯度吃掉 14 GB
- Master weights 吃掉 28 GB
- Adam 状态吃掉 56 GB
- Activations 吃掉 30 GB
合计 128 GB 的额外开销。Model States 就有 112 GB，是权重本身的 8 倍。
“模型参数 14 GB"描述的只是冰山一角。训练时的完整显存由 Model States + Activations 共同决定，参数只是其中一个分项。
Activation Checkpointing 能省多少#
之前写过，activation checkpointing 的本质是缩短中间 tensor 的存活时间：只保留每层的边界 activation，backward 到达时重新跑一遍 forward 来恢复中间结果。概念不重复了，这里只算数字。
如果采用"每层只保留输入"的 full checkpointing：
$$ M_{\text{act, ckpt}} = L \times BSd \times 2 = 32 \times 1 \times 2048 \times 4096 \times 2 \approx 0.5\ \text{GB} $$从 30.6 GB 降到 0.5 GB，activation 显存减少了 98%。代价是额外约 33% 的 FLOPs（backward 中需要重新执行一遍 forward）。
但问题是：Activation checkpointing 不碰 Model States。
$$ 112\ \text{GB (Model States)} + 0.5\ \text{GB (Activations)} + 2\ \text{GB (overhead)} \approx 114.5\ \text{GB} $$还是 80 GB 的 1.4 倍。即使把 Activations 砍到接近零，单卡仍然放不下 7B 模型的 Model States。
Activation checkpointing 解决的是"Activations 太大"的问题。但 7B 训练的瓶颈不在 Activations——在 Model States。
Offloading：把 GPU 放不下的数据搬到更便宜的地方#
Activation checkpointing 用重新计算换显存。Offloading 用另一种资源换——数据搬运。
思路#
GPU 的 HBM 只有 80 GB，但主机的 CPU 内存通常有几百 GB 甚至上 TB。112 GB 的 Model States 中，有很大一部分在 forward/backward 期间并不被使用：
- Master weights 和 Adam \(m\)、\(v\)：只在 optimizer step 时才需要
- Forward 和 backward 阶段，这 84 GB 的 FP32 数据只是占着 GPU 显存，什么也没做
CPU offload 的思路就是：把这些"暂时用不到"的数据搬到 CPU 内存，需要的时候再搬回来。
CPU Offload#
典型做法（DeepSpeed ZeRO-Offload 的思路）：
- Forward/backward 阶段：GPU 上只保留 BF16 参数（14 GB）和 BF16 梯度（14 GB）
- Optimizer step：把梯度发给 CPU，CPU 上用 FP32 master weights 和 Adam 状态完成参数更新，再把更新后的 BF16 参数发回 GPU
这样 GPU 上的常驻显存：
$$ 14\ \text{(params)} + 14\ \text{(grads)} + 0.5\ \text{(ckpt activations)} + 2\ \text{(overhead)} \approx 30.5\ \text{GB} $$CPU 内存承担 84 GB（master weights + Adam 状态）。
代价#
Offloading 不是免费的。GPU 和 CPU 之间通过 PCIe 连接，带宽有限：
- PCIe 4.0 x16 的理论带宽约 32 GB/s（单向）
- 每个 optimizer step 需要传输梯度到 CPU + 更新后的参数回 GPU
- 84 GB 数据的搬运在 PCIe 带宽下需要数秒
如果 forward + backward 的计算时间不够长，GPU 会空等 CPU 完成 optimizer step 并回传参数。计算量越小（batch size 越小），offload 的性能损失越明显。
NVMe Offload#
更极端的做法：连 CPU 内存也不够时，还可以把数据放到 NVMe SSD 上。
好处是容量几乎不受限。代价是 NVMe 的带宽比 PCIe 更低（典型 3–7 GB/s），搬运延迟更大。
两种省显存策略的对比#
| Activation Checkpointing | Offloading | |
|---|---|---|
| 用什么换显存 | 重新计算 | 数据搬运 |
| 省的是哪部分 | Activations | Model States（optimizer 等） |
| 代价 | 额外 FLOPs（约 33%） | PCIe/NVMe 带宽和延迟 |
| 能否独立使用 | 能，但不解决 Model States | 能，但不解决 Activations |
两者解决的是不同的问题，通常需要组合使用。 对 7B 模型来说，activation checkpointing 把 Activations 从 30 GB 砍到 0.5 GB，offloading 把 84 GB 的 optimizer 相关数据搬到 CPU——两个一起用，GPU 上的显存才降到 80 GB 以内。
写在最后#
回到标题：7B 模型只有 14 GB 权重，一张 80 GB GPU 为什么仍可能 OOM？
因为训练显存不等于权重显存。 BF16 mixed precision + AdamW 下，每个参数对应 16 bytes 的 Model States——7B 模型的 Model States 就有 112 GB，仅此一项已经超过 80 GB。再加上约 30 GB 的 Activations，单卡训练至少需要 140+ GB。
要把 7B 训练装进单张 80 GB GPU，至少需要同时使用 activation checkpointing（用计算换 activation 显存）和 optimizer offloading（用 PCIe 带宽换 model states 显存）。
但这篇只算了单卡的账。真正训练 7B 甚至更大的模型，更常见的做法是把这张放不下的账本拆到多张卡上——ZeRO / FSDP 把 Model States 分片，Tensor Parallelism 把单层的计算和 activation 拆开，Pipeline Parallelism 把不同的层放到不同的卡上。这些是下一篇的内容。
参考资料#
- 从参数量到 Peak Memory：模型显存到底应该怎么算？
- 从推理到训练：模型调度中的显存约束
- 从 AMP 到 Mixed Precision：训练和推理为什么不使用同一种精度
- 从 Forward 到 Optimizer：再看一遍 Training Loop
- MIT 6.S191: Introduction to Deep Learning — Lecture 9: Secrets to Massively Parallel Training. https://introtodeeplearning.com/
- Korthikanti et al. Reducing Activation Recomputation in Large Transformer Models. 2022.（Activation 显存公式的来源）
- Rajbhandari et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC 2020.
- Ren et al. ZeRO-Offload: Democratizing Billion-Scale Model Training. USENIX ATC 2021.