激活内存：重计算与卸载

并行切分解决的是"模型装不下"，这一章解决"激活装不下"。大 batch、长序列下，前向激活的大小能超过权重本身——7B 模型长序列训练时激活可达数百 GB。两条路：重计算（不存激活，反向时重算）与卸载（搬到 CPU 内存）。两者都是"用时间换显存"，关键是知道各自的代价与适用场景。

激活有多大

先算清楚。Transformer 每层激活的字节数（16 位、无重计算、不含 attention 矩阵）：

\text{每层激活} \approx 20 \sim 34 \times B \times s \times h \quad \text{字节}

（系数取决于实现，含 QKV、注意力输出、FFN 中间层、LayerNorm 等。）

以 7B 模型（ $h=4096$ ）、batch 8、seqlen 4096、32 层：

32 \text{层} \times 24 \times 8 \times 4096 \times 4096 \times 2\ \text{B} \approx 200\ \text{GB}

激活 200 GB，权重才 14 GB——大 batch/长序列下激活是显存的第一大户。FP8/BF16 注意力矩阵与 KV cache 还要另算。

重计算：时间换显存

重计算（recomputation/checkpointing）：前向时不保存激活（或只存少量检查点），反向时需要的地方重算一遍。

策略	保存什么	激活显存	额外计算
全保存（无重计算）	全部激活	100%	0
选择性重计算	只存便宜算子的输入（如 attention 的 QKV 输入）	~30–40%	~1–3%
每层重计算（block 级 checkpoint）	每层只存输入	~ $1/L$	+33%（一次额外前向）
完全重计算	只存最初输入	极小	+100%

实践结论：

每层重计算是默认起点：显存除以层数，代价是一次额外前向（反向 FLOPs ≈ 2× 前向，共 3 次前向的计算量，+33%）
选择性重计算是优化的甜蜜点：attention softmax 的中间量又大又便宜（重算快），保留它们、重算贵的（矩阵乘输出），显存 ~1/3、计算只多 1–3%
FlashAttention 本身就是重计算：不存 $s^2$ 的注意力矩阵，反向时重算——它证明了"重算便宜的部分"几乎免费

判断标准：重算一个算子的成本 = 它的 FLOPs ÷ 带宽搬运的收益。逐元素/softmax 类算子重算近乎免费（访存受限），矩阵乘重算很贵（计算受限）——选择性重计算就是只重算前者。

卸载：DRAM 换显存

卸载（offload，如 DeepSpeed ZeRO-Offload/Infinity）把激活（或优化器状态）搬到 CPU 内存，需要时搬回：

卸载什么	显存收益	代价	前提
优化器状态	12 B/参数 → 0	每 step 两次 PCIe 搬运（ZeRO-1 分片下很小）	PCIe 带宽
激活	大（batch×seqlen 级）	前向搬出 + 反向搬回，PCIe 64 GB/s	激活越大越亏
参数（ZeRO-Infinity）	2 B/参数	每层 all-gather 走 NVMe/PCIe	很慢，末选

关键算术：PCIe 5.0 x16 = 64 GB/s，而 HBM 是 3.35 TB/s——差 50 倍。卸载的代价公式：

T_{\text{offload}} = \frac{\text{搬运量} \times 2}{64\ \text{GB/s}} \qquad \text{vs} \qquad T_{\text{recompute}} = \text{重算 FLOPs} \div \text{算力}

以 7B、batch 8、seqlen 4096 的激活（~25 GB，每层卸载）：每层 50 GB 往返 ÷ 64 GB/s ≈ 0.8 ms/层，32 层 ~25 ms/step——对 step time 1 s 的训练可接受；但 batch 再大、序列再长，搬运时间线性涨，很快超过重计算的 33%。

经验法则：

先用选择性重计算（几乎免费）
不够再加每层重计算（+33%）
还不够卸载优化器状态（ZeRO-Offload，代价小）
最后才卸载激活或参数（代价大，等于宣布 PCIe 是新瓶颈）

NVMe 与"更远的内存"

ZeRO-Infinity 把卸载推到 NVMe（~10 GB/s，容量 TB 级）：显存 → CPU 内存 → NVMe 三级。每往下一级，带宽掉 5–6 倍——与 GPU 内存层次同构（见内存层次），只是这次由软件管理。适用场景是"显存与 CPU 内存都不够的超大模型实验"，生产训练很少走到这一级。

与并行策略的组合

激活技术与并行正交、且互相影响：

技术	显存收益	时间代价	与并行的关系
选择性重计算	激活 ÷3	+1–3%	独立生效
每层重计算	激活 ÷L	+33%	独立生效
优化器卸载	12 B/参数	小	配合 ZeRO-1
激活卸载	大	线性涨	与 TP 冲突（TP 已省激活，重复省）

注意最后一行：TP=8 已经把激活除了 8，再卸载激活的收益要除以 TP 度重新算——先并行、后重计算、最后卸载，这个顺序里每一步都重新核算。

本章要点

激活大小 ≈ (20–34) × batch × seqlen × hidden × 层数字节；大 batch/长序列下是显存第一大户。
重计算阶梯：选择性（+1–3%，显存 ÷3）→ 每层（+33%，显存 ÷L）→ 完全（+100%，末选）。
重算便宜的是访存受限算子（softmax/逐元素），贵的是计算受限算子（矩阵乘）——选择性重计算只重算前者。
卸载按 PCIe 64 GB/s 计价，比 HBM 慢 50 倍；优化器卸载便宜，激活/参数卸载昂贵。
顺序：先并行切分、再选择性重计算、再每层重计算、再优化器卸载、最后才激活/参数卸载。