> For the complete documentation index, see [llms.txt](https://yeasy.gitbook.io/llm_internals/llms.txt). Markdown versions of documentation pages are available by appending `.md` to page URLs; this page is available as [Markdown](https://yeasy.gitbook.io/llm_internals/di-er-bu-fen-xun-lian-pian/07_distributed_training/7.5_activation_checkpointing.md).

# 7.5 激活重计算：用时间换空间

前四节分掉的都是模型状态：参数、梯度和优化器状态。训练时显存里还有另一大块，即前向传播留给反向使用的**激活值**（activations）。7.4.2 说明 1F1B 只限制了同时持有几份激活，每一份有多大还没有处理。本节先逐项清点一层要存多少激活，再讲三种安置办法：留在显存、丢掉重算、卸载到主机内存，最后说明它们怎样影响效率指标 MFU 与 HFU。

## 7.5.1 一层要存多少激活：逐项清点

**要解决的问题。** 反向传播计算某个算子的梯度时，要用到它前向的输入，有时还要用到输出。前向因此不能算完就丢，每层都要留下一批张量。

[Megatron 的论文](https://arxiv.org/abs/2205.05198)对标准 Transformer 层（GeLU 的 MLP、中间维 `4h`、带 dropout）做了逐项清点。记序列长度 $s$、微批量 $b$、隐藏维 $h$、头数 $a$，依次对应 3.8 节表 3-26 的 `T`、`B`、`d_model`、`n_h`；激活按 16 位存放，每个元素 2 字节，dropout 掩码每个元素 1 字节。

| 位置  | 要保存的张量                     |               字节数 |
| --- | -------------------------- | ----------------: |
| 注意力 | Q/K/V 投影的共同输入              |            $2sbh$ |
| 注意力 | Q 与 K（供 $QK^{T}$ 反向）       |            $4sbh$ |
| 注意力 | Softmax 输出 `[a, s, s]`     |        $2as^{2}b$ |
| 注意力 | Softmax 后的 dropout 掩码      |         $as^{2}b$ |
| 注意力 | dropout 的输出与 V（供“权重乘 V”反向） | $2as^{2}b + 2sbh$ |
| 注意力 | 输出投影的输入，及其后的 dropout 掩码    |      $2sbh + sbh$ |
| MLP | 两个线性层的输入                   |     $2sbh + 8sbh$ |
| MLP | GeLU 的输入，及末尾的 dropout 掩码   |      $8sbh + sbh$ |
| 归一化 | 两个 LayerNorm 的输入           |            $4sbh$ |

表 7-14：一层 Transformer 为反向保存的激活，取自 Megatron 论文第 4.1 节。注意力合计 $11sbh + 5as^{2}b$，MLP 合计 $19sbh$，归一化 $4sbh$。

三项相加：

$$
\text{每层激活} = sbh\left(34 + 5,\frac{as}{h}\right)\ \text{字节}
$$

括号里的 34 随序列长度线性增长，$5as/h$ 一项来自 `[s, s]` 的注意力矩阵，随序列长度平方增长。

**代入 GPT-3 175B。** `s = 2048`、`h = 12288`、`a = 96`、`b = 1`。线性项 `34 × 2048 × 12288 ≈ 0.856 GB`；平方项 `5 × 96 × 2048² ≈ 2.013 GB`，相当于 `5as/h = 80`。每层合计 2.87 GB，其中平方项占 70%。96 层共 275 GB，而这个模型的 BF16 参数是 350 GB。一个微批量、2K 的序列，激活已接近参数的规模；批量加大或序列变长，激活反超参数。

**FlashAttention 改变了这张表。** [10.3 节](/llm_internals/di-san-bu-fen-tui-li-yu-bu-shu-pian/10_inference_optimization/10.3_flash_attention.md)的 FlashAttention 前向不保存 `[s, s]` 的矩阵，只保存注意力输出和每行一个 LSE，反向时分块重算。用了它，表中三个含 $as^{2}b$ 的项就消失了。它本身就是一种内核内部的重计算。SwiGLU、无 dropout、GQA 的模型也要按同样的方法重新清点，系数不再是 34，方法不变。

## 7.5.2 重计算的机制

**激活重计算**（activation recomputation）也称 activation checkpointing 或 gradient checkpointing。它放弃保存一部分前向激活，反向用到时再算一遍：

1. **前向。** 把模型划成若干段，只保存每段的输入，称为检查点；段内的中间激活算完即丢。
2. **反向。** 反向传播走到某一段时，从该段的检查点出发，重做一次这一段的前向，得到段内激活，算出梯度后再丢弃。

**重算必须复现同一次前向。** dropout 依赖随机数，重算时若掩码不同，梯度就是错的。PyTorch 的 `torch.utils.checkpoint.checkpoint` 默认 `preserve_rng_state=True`，在检查点处保存并恢复随机数状态；[文档](https://docs.pytorch.org/docs/stable/checkpoint.html)建议使用 `use_reentrant=False` 的实现。张量并行下还要区分两类随机数状态：Megatron-LM 的 `CudaRNGStatesTracker` 让张量并行区域内的 dropout 在各卡上取不同的种子，区域外的 dropout 取相同的种子。

## 7.5.3 代价怎么算：三分之一的来历与分段数的选择

**计算的代价。** 一层的反向约为前向的 2 倍计算量，因为要分别对输入和对权重求梯度（Megatron 论文附录 A）。不重算时一步是 `1 + 2 = 3` 份前向的量；每层都重算一次前向，变成 `1 + 2 + 1 = 4` 份，增加三分之一。Megatron 论文的实测与此相符：全量重计算带来 30% 到 40% 的时间开销，530B 模型上是 36%。

**分段数怎么选。** 设共 $L$ 层，分成 $k$ 段，每个检查点大小为 $c$，一层的完整激活为 $A$。反向处理某一段时，显存里有 $k$ 个检查点，加上这一段 $L/k$ 层的完整激活：

$$
\text{峰值} \approx k,c + \frac{L}{k},A
\qquad\Rightarrow\qquad
k^{\*} = \sqrt{\frac{L,A}{c}}
$$

[Chen 等人的经典结果](https://arxiv.org/abs/1604.06174)讨论的是 $A = c$ 的均匀链式网络，此时 $k^{\*} = \sqrt{L}$，显存从 $O(L)$ 降到 $O(\sqrt{L})$，代价是多一次前向。把检查点压到 $O(1)$ 个并不能让显存也降到 $O(1)$：段内的 $L/k$ 项随之变成 $O(L)$。要继续降只能递归地分段，同一篇论文给出的极限是 $O(\log L)$ 的显存、$O(L\log L)$ 的前向计算。

Transformer 不是均匀链。GPT-3 175B 的检查点是层输入 `2sbh ≈ 0.050 GB`，一层的完整激活是 2.87 GB，相差 57 倍：

|           分段数 $k$ | 检查点 $k,c$ | 段内激活 $(L/k)A$ |       峰值 |
| ----------------: | --------: | ------------: | -------: |
|                 1 |   0.05 GB |      275.4 GB | 275.5 GB |
| 10（约 $\sqrt{96}$） |   0.50 GB |       27.5 GB |  28.0 GB |
|      74（$k^{\*}$） |   3.72 GB |       3.72 GB |   7.4 GB |
|          96（每层一个） |   4.83 GB |       2.87 GB |   7.7 GB |

表 7-15：GPT-3 175B 在不同分段数下的激活峰值，`L = 96`、`b = 1`、`s = 2048`。

最优分段数是 74，与“每层一个检查点”几乎没有差别，而 $\sqrt{L} \approx 10$ 的峰值是它的近 4 倍。实际框架的全量重计算因此以层为单位：每层只存输入 $2sbh$，总计 $2sbhL$，反向时一次只展开一层。

## 7.5.4 选择性重计算

全量重计算把 275 GB 压到不足 8 GB，代价是 33% 的计算。Megatron 论文提出**选择性激活重计算**（selective activation recomputation）：只重算显存占用大而计算量小的部分。

符合这一条件的正是 7.5.1 中 $5as/h$ 对应的几步：$QK^{T}$、Softmax、Softmax 后的 dropout、权重乘 V。它们的张量是 `[a, s, s]`，体积大；每个元素上的运算却很少，没有大矩阵乘法。只丢弃并重算这一部分：

* **省下的显存。** GPT-3 175B 每层从 2.87 GB 降到 0.856 GB，省 70%；MT-NLG 530B 的 `5as/h = 64`，省 65%。
* **多出的计算。** 只需重做注意力矩阵的计算和权重乘 V 两步。论文附录 A 推得硬件 FLOPs 与模型 FLOPs 之比约为 $1 + s/6h$。代入 GPT-3 175B 是 `1 + 2048 ÷ 73,728 ≈ 1.028`，计入词表项后为 1.027，论文报告的开销是 2.7%，对照全量重计算的 33%。

表 7-16 汇总各种组合下的每层激活，张量并行与序列并行的部分见 7.3.6。

| 配置                   | 每层激活（字节）                    | GPT-3 175B、`t = 8` 时 |
| -------------------- | --------------------------- | -------------------: |
| 不并行                  | $sbh(34 + 5as/h)$           |              2.87 GB |
| 张量并行                 | $sbh(10 + 24/t + 5as/(ht))$ |              0.58 GB |
| 张量并行 + 序列并行          | $sbh(34/t + 5as/(ht))$      |              0.36 GB |
| 张量并行 + 序列并行 + 选择性重计算 | $sbh \cdot 34/t$            |             0.107 GB |
| 全量重计算                | $2sbh$                      |             0.050 GB |

表 7-16：各种技术组合下的每层激活，公式取自 Megatron 论文表 2。

三者叠加后每层 0.107 GB，96 层约 10.3 GB，已与全量重计算的 4.8 GB 同一量级，计算开销却只有 2.7%。论文在 530B 模型上的端到端结果是吞吐比全量重计算高 29.7%。

**与 FlashAttention 的关系。** 选择性重计算丢掉的那几项，正是 FlashAttention 本来就不保存的。用了 FlashAttention，这部分收益已经在内核里兑现，选择性重计算只剩下其余可选模块。

**实现锚点。** Megatron-LM 用 `recompute_granularity` 选粒度：`selective` 重算 `recompute_modules` 指定的子模块，默认是 `core_attn`；`full` 重算整层，此时由 `recompute_method` 决定哪些层参与，`uniform` 是把各层均匀分组、每组存一个检查点，`block` 是每个流水线级只对指定数量的层重算、其余层照常保存。后者用于显存还有余量时少做一些重算。

## 7.5.5 第三条路：把激活值卸载到主机内存

前面几小节把权衡讲成了两端：要么存显存，要么花算力重算。还有第三种做法：把选中的激活在前向阶段搬到主机内存，反向需要时再传回来。它既不占显存，也不重算，付出的是互连带宽。

它变得实用，是因为硬件条件变了。NVIDIA 的一篇博客说明了这一点，题为 [Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading](https://developer.nvidia.com/blog/reducing-high-bandwidth-memory-bottlenecks-in-jax-based-llm-training-with-host-offloading/)，发布于 2026-07-10。Grace CPU 与 Blackwell GPU 之间通过 NVLink-C2C 相连，双向带宽 900 GB/s，锁页主机内存因此成为可用的中转区。下一代 Vera Rubin 把这一数字提高到 1.8 TB/s。在带宽低得多的 PCIe 链路上，这条路的性价比要差得多。

激活的安置于是成为三选一。这是本书的归纳，NVIDIA 原文并未如此表述。

| 策略                 | 省下什么     | 付出什么           |
| ------------------ | -------- | -------------- |
| 保存在显存              | 不重算、不搬运  | 显存             |
| 重计算（7.5.2 到 7.5.4） | 显存       | 算力             |
| 卸载到主机内存            | 显存与重算的算力 | 互连带宽，以及能否与计算重叠 |

表 7-17：激活值的三种安置策略。

最后一列的后半句最容易被忽略：有带宽还不够，还要有独立的计算把传输藏在后面。同一篇文章给出了实测数据，配置是 DeepSeek-V3 671B、GB200 NVL72、128 卡。未经调优的主机卸载比重计算慢约 6.3%：541.6 对 578.3 TFLOPs/s/device，`1 − 541.6 ÷ 578.3`。启用延迟隐藏调度与流水化传输之后，才反超到 908.2 TFLOPs/s/device，比重计算快约 57%。

所以不能简化成“卸载比重计算更快”。准确的说法是：互连带宽充足，且有足够的独立计算可供重叠时，卸载才优于重计算；否则它只是把显存压力换成一段暴露在关键路径上的传输。卸载也不总是降低显存峰值。在上面的最优配置里，因为要为拷贝缓冲和预取的激活留空间，GPU 峰值显存从重计算的 151.3 GiB 升到 165.2 GiB。它真正改变的是相对于“把激活直接留在显存”而言哪些批量配置变得可行：后者在同一配置下直接 OOM，而重计算也能跑通，只是慢得多。

## 7.5.6 怎么报效率：MFU 与 HFU

上一小节的 TFLOPs/s/device 是绝对值，跨硬件不可比。用来横向比较的是利用率：实测吞吐除以硬件峰值。分子算哪些浮点运算，有两种口径：

* **模型 FLOPs**：完成一次前向加反向在数学上必需的浮点运算，只由模型架构决定，与实现和硬件无关，即 [3.7 节](/llm_internals/di-yi-bu-fen-ji-chu-pian/03_components/3.7_full_architecture.md)（3.7.7）的 $6ND$。
* **硬件 FLOPs**：硬件上实际执行的浮点运算。启用了激活重计算，多做的前向也算在内，因此硬件 FLOPs 大于模型 FLOPs。

两者分别除以“迭代时间 × 加速器峰值算力”，得到 **MFU**（Model FLOPs Utilization）与 **HFU**（Hardware FLOPs Utilization）。[Megatron 的论文](https://arxiv.org/abs/2205.05198)采用这一定义，7.5.4 的 $1 + s/6h$ 就是两者之比。

**手算一次 MFU。** Megatron 论文表 3 与表 5 给出 175B 模型的评测配置：64 张 A100，全局批量 64 条、每条 2,048 词元，每步 13.75 s。每步 `64 × 2048 = 131,072` 个词元：

* 按 $6ND$：`6 × 1,750 亿 × 131,072 ≈ 1.376 × 10¹⁷` 次。
* 每卡每秒：`1.376 × 10¹⁷ ÷ 13.75 ÷ 64 ≈ 156 TFLOPs/s`。
* 除以 A100 的峰值 312 TFLOPs/s，得 50.1%。

论文用的是含注意力项和词表项的精确公式，模型 FLOPs 为 `1.411 × 10¹⁷`，得 51.4%，与其表 5 一致；换成硬件 FLOPs 的公式得 52.8%，即 HFU。$6ND$ 略去了注意力的平方项，所以偏低约 1 个百分点。[PaLM 论文](https://arxiv.org/abs/2204.02311)附录 B 给出同样的算法：`238.3K 词元/s × 6 × 5,400 亿 ÷ (275 TFLOPs/s × 6,144)`，得 45.7%，计入注意力后为 46.2%。

**开启重计算会让 HFU 上升、MFU 下降。** 重计算增加了硬件实际执行的运算，HFU 的分子变大；它没有让模型多处理任何词元，只延长了迭代时间，MFU 下降。PaLM 论文提出 MFU 的动机就在这里。它指出 HFU 不是一致而有意义的训练效率指标。MFU 的理论最大吞吐只计入前向与反向必需的运算，不计重物化（rematerialization），不同实现之间才能公平比较。PaLM 540B 报告的是一对数字：MFU 46.2%，HFU 57.8%。

看到一个利用率数字时，有三件事要问清楚：

1. **报的是 MFU 还是 HFU。** 同一次训练两个数字可以差十几个百分点，PaLM 的 46.2 与 57.8 就是例子。
2. **分母用的哪个峰值。** 同一张卡的 BF16 稠密峰值、FP8 峰值、带结构化稀疏的峰值可以相差一倍以上，换个分母 MFU 就能凭空翻倍。Megatron 的论文在脚注里写明它用的 A100 峰值是 312 TFLOPs/s，Llama 3 的表 4 注明是 BF16 MFU，都是值得效仿的写法。
3. **有没有把优化效果算进去。** Megatron 在 2,240 张 A100 上训练 530B 模型时，MFU 从全量重计算的 42.1% 提到 54.2%。提升来自少做无用功，不是模型变好了。

重计算和卸载都是拿 MFU 换显存：它们让原本 OOM 的配置跑得起来，代价是同样的时间里模型 FLOPs 变少。7.5.5 那组 541.6、578.3、908.2 的对比，就是在同一块显存约束下比较几种换法各剩多少 MFU。这些办法调的都是“哪些数留下”，每个数自身占几位还没有动。