动画图解 LLM 系统 · 激活重算

Gradient checkpointing

训练时，用多算一遍前向，换回被激活值占满的显存

按 → 或点击推进动画 · 页面底部的进度条可以拖动回看

目录

目录

  1. 反向为什么需要激活值
  2. 显存锯齿：前向堆、反向拆
  3. 一层 Transformer 存了多少字节
  4. checkpoint 的思路与逐层重算
  5. 代价：多一次前向 ≈ +33%
  6. 分段与 √n 策略（Chen et al. 2016）
  7. 交互：每 k 层一个 checkpoint
  8. 选择性重算：只重算 attention core
  9. FlashAttention：kernel 里的重算
  10. CPU offload
  11. 与流水线并行、FSDP 的配合
  12. torch.utils.checkpoint 与按算子选择的策略
动机

反向要用前向的输入

一个线性层 \(y = Wx\)。前向算完 \(y\) 就往下传，但 \(x\) 不能扔：它被存进 autograd 的"保存区"。

反向时 \(dy\) 回来：\(dx = W^\top dy\) 只要权重；\(dW = dy\,x^\top\) 却必须用到前向的 \(x\)。

非线性也一样：GeLU、softmax、归一化的导数都依赖它们的输入或输出。这些在前向产生、留给反向用的张量，就叫激活（activations）。

显存锯齿

前向一层层堆，反向一层层拆

前向经过每一层，都把这一层的激活压进显存（右侧的栈），左边的曲线是显存随时间的变化。

反向从最后一层往回走，用完一层的激活就释放一层：后进先出。反向每层的计算量约是前向的 2 倍，所以下降段更长。

每个训练 step 重复一次：显存曲线是锯齿形，只看图中保存的激活，峰值在前向结束时：各层保存量相加。

账本

一层 Transformer 存了多少

按标准 attention 的 Transformer 账本数一遍（16 位激活，\(s\) 序列长、\(b\) micro-batch、\(h\) 隐藏维、\(a\) 头数）：每层 \(sbh\,(34 + 5as/h)\) 字节。图中按 GPT-3 175B 的形状画（\(a=96,\ s=2048,\ h=12288\)，\(5as/h = 80\)）。

\(b=1\) 时一层约 2.87 GB，96 层就是 275 GB：远超一张 80 GB 的卡，而这还没算权重和优化器状态。

和 \(s^2\) 成正比的那一块（softmax 输出、dropout mask 与输出）占了 70%。FlashAttention 不存这个 \(s \times s\) 矩阵，剩下 \(34\,sbh\)：一层 0.86 GB，96 层仍有 82 GB。

显存账本 · 先区分哪些能省

checkpoint 主要改变激活的保存方式

峰值 ≈ 参数 + 梯度 + 优化器状态
　　 + 活跃激活 + 临时工作区 + 通信缓冲

适合解决的问题

长序列、大 micro-batch、多层网络让活跃激活变大；重算减少为反向长期保留的中间结果。

换策略也省不掉的部分

参数、优化器状态不会因为包上 checkpoint 就自动分片。若它们已占满显存，需要结合低精度、ZeRO/FSDP 或减少模型规模。

公式中的 A、c 必须来自当前算子与精度实现。FlashAttention、SwiGLU、融合算子与 TP 都会改变激活账本，不能把 34sbh 当成所有 Transformer 的常数。

思路

只存层的输入，反向前重算一遍

正常训练：一层里的每个算子都留下激活，全部要等到反向才释放。

checkpoint：只把这一层的输入存下来（青色）；层内的中间结果算完下一步就丢，前向时显存几乎不涨。

轮到这一层反向时，先拿存好的输入把这层的前向重跑一遍，恢复全部激活，再照常反向、用完释放。

逐层重算

每层一个 checkpoint

前向：每层只留一个输入（青色细条），层内激活即用即丢。虚线是不重算时的显存曲线。

反向：每到一层，先重算它（橙色），再反向释放。任何时刻最多只有一层的完整激活。

峰值 ≈ \(L\) 个层输入 + 1 层完整激活。GPT-3 175B 形状下：\(96 \times 2sbh \approx 4.8\) GB，加一层约 2.9 GB，从 275 GB 降到不到 8 GB。代价是时间变长了。

演算 · 四层标量链

重算恢复数值，梯度仍沿原链返回

教学模型 x₀ = 1，xᵢ₊₁ = 2xᵢ，共 4 层，损失 L = x₄。每两层一段，保留入口 x₀ = 1、x₂ = 4。

先重算后半段 4 → 8 → 16，反传梯度 1 → 2 → 4；再从 x₀ 重算前半段，继续得到 dL/dx₀ = 16。

代价

多一次前向：约 +33%

一层的反向要算 \(dx\) 和 \(dW\) 两个矩阵乘，约是前向的 2 倍。不重算时，一个 step 的计算量是 1 + 2 = 3 份。

全部重算后，每层反向前多一次前向：4 份，计算量 +33%。实测常常更多（36%–39%），因为重算还会打乱原本和计算重叠的通信。

时间账本 · “多 33%”需要前提

实际 step 还包含通信和优化器

T₀ = F + B + C
T₁ ≈ F + B + C + R + ΔC

F 是前向，B 是反向，C 是其他关键路径时间；R 是重算，ΔC 是重算改变重叠后新增或减少的关键路径时间。

教学例子耗时
F=20 ms，B=40 ms，C=10 ms原 step：70 ms
全重算 R=20 ms，ΔC=0新 step：90 ms，增加 28.6%
再多 5 ms 无法隐藏的通信新 step：95 ms，增加 35.7%

只有 B≈2F、C 可忽略、R≈F 且重叠方式不变时，才得到 4F/3F≈1.33。比较训练吞吐时还应保持每 step 的 token 数可比。

分段

每 k 层存一次：显存 n/k + k

16 层分成 4 段，每段 \(k = 4\) 层。前向只在段的入口存 checkpoint，整段内部都不存。

反向时一次重算一整段，段内 \(k\) 层的激活同时在显存里，再逐层反向释放。重算总量仍然是一次完整的前向。

峰值 ≈ \(n/k\) 个 checkpoint + \(k\) 层激活：\(k\) 太小，checkpoint 多；\(k\) 太大，重算的段大。中间有一个最优。

推导

√n 从哪里来

设每层激活与 checkpoint 一样大（都记 1 份）：$$M(k) = \frac{n}{k} + k$$
$$M'(k) = 1-\frac{n}{k^2} = 0 \;\Rightarrow\; k = \sqrt n$$ 此时 \(M_{\min} = 2\sqrt n\)。
额外计算仍是一次前向，显存从 \(O(n)\) 变成 \(O(\sqrt n)\)。递归地把每段再分段，可以降到 \(O(\log n)\)，代价是更多次重算。
Transformer 里 checkpoint 只是层的输入 \(c = 2sbh\)，远小于一层激活 \(A = 34sbh\)：$$M = \frac{n}{k}c + kA,\quad k^* = \sqrt{nc/A}$$ \(n = 96\) 时 \(k^* \approx 2.4\)，这份估算下先比较每 2 层或每 3 层一段，再把工作区和真实峰值纳入测量。

蓝线是 checkpoint 的份数 \(n/k\)，橙线是重算一段时的 \(k\) 层激活，黑线是两者之和（\(n = 16\)）。

最小值在 \(k = \sqrt n\)。这个策略来自 Chen et al. 2016：1000 层的 ResNet 显存从 48 GB 降到 7 GB，耗时增加约 30%。

演算 · 有单位的分段预算

检查邻近整数，比照搬 √n 更可靠

教学模型：48 层，每段入口 c = 16 MiB，每层重算时激活 A = 128 MiB。先忽略参数、工作区与释放时序。

M(k) ≈ ceil(48/k)×16 + 128k MiB。连续近似 k* = √6 ≈ 2.45；整数 k=2、3 都是 640 MiB。

推导 · 这个最优值依赖什么

一个简化模型，不能替代显存时间线

M(k) ≈ (n/k)c + kA
dM/dk = −nc/k² + A = 0 → k* = √(nc/A)

模型假设

层大小相同、顺序链结构、每段完整重算一次；边界各保存一份，反向段内激活最多保留 k 层。

实际峰值还会移动

边界会逐段释放，末段可能不重算，non-reentrant 可提前停止；workspace、通信与不均匀层会叠到不同时间点。

先用模型找候选 k，再用真实 step 测峰值。选 k 时还要满足显存余量和吞吐，而不是只追求最小的激活估计。

交互

显存与时间的权衡

激活峰值
不重算时
额外计算
显存节省

示意模型：GPT-3 175B 形状、\(b=1\)、FlashAttention、单卡不切 TP，每层激活 \(34sbh \approx 0.86\) GB，checkpoint \(2sbh \approx 0.05\) GB；前 N 层每 k 层存一个 checkpoint，其余层不重算。前向 1 份时间、重算 1 份、反向 2 份；不计重算时一层的临时工作区。灰色虚线是不重算的曲线。

选择性重算

只重算最占显存、最便宜的那部分

同一层里，不同激活"每字节的重算代价"差别很大。\(QK^\top\)、softmax、dropout、乘 \(V\) 这几步（attention core）产生的激活和 \(s^2\) 成正比；在这里的 GPT-3 形状下，它占总 FLOPs 的比例很小。

只把这几步重算，其余照存：每层激活从 \(sbh(34 + 5as/h)\) 降到 \(34\,sbh\)。GPT-3 省 70%，MT-NLG 省 65%。

代价只有 2.7%（GPT-3）和 1.6%（MT-NLG）的额外 FLOPs，远低于整层重算的 33%。再配合序列并行，把剩下的 \(34\,sbh\) 也按张量并行度 \(t\) 切开。

策略 · 按“省多少 / 多算多久”选择

不同层和算子的价格不一样

教学测量：可重算区 A 省 400 MiB、增加 2 ms；B 省 100 MiB、增加 5 ms；C 省 300 MiB、增加 3 ms。目标至少腾出 600 MiB。

选择估计节省额外时间
A + B500 MiB7 ms：容量仍不足
A + C700 MiB5 ms：满足目标
B + C400 MiB8 ms：容量仍不足

A 的单位时间收益是 200 MiB/ms，C 是 100，B 是 20。先重算“占得多、算得便宜”的区域通常更划算。

这是独立区域、节省可相加的教学模型。真实峰值由张量生命周期决定，两个区域若不同时存活，峰值收益未必相加。

实测

一层的前向 + 反向时间（22B 模型）

22B 模型一层的实测：整层重算让反向从 11.9 ms 涨到 19.5 ms，总时间 +39%。

选择性重算只多 7%；再加序列并行（LayerNorm、dropout 也按 \(t\) 切开），总开销只剩 4%。模型越大越明显：530B 和 1T 上从 36% 降到 2%，530B 的 MFU 从 42.1% 提到 54.2%。

FlashAttention

重算已经藏在 attention kernel 里

标准实现把 \(s\times s\) 的注意力概率 \(P\) 整个写进显存留给反向。FlashAttention 按块算：每块分数算完、累加进输出就丢，保留 Q、K、V 和输出 \(O\)，另外存每行一个 logsumexp。

反向时用 \(Q, K, V\) 和 logsumexp 把每块 \(P\) 重算出来。这正是选择性重算，只是做在一个 kernel 里，不用读写 HBM 上的 \(P\)，往往比不重算还快。

CPU offload

另一条路：把激活搬到 CPU 内存

不重算，而是在前向时把激活经 PCIe 异步拷到主机内存（D2H），反向前再提前拷回（H2D），拷贝走单独的 stream，和计算重叠。

问题在带宽：H100 的 HBM 是 3.35 TB/s，PCIe Gen5 x16 单向约 64 GB/s，差 50 倍。拷贝比计算慢时，反向就得停下来等（红色）。一台机器 8 张卡还要共享主机的内存带宽。

PyTorch 里一行就能试：with torch.autograd.graph.save_on_cpu(pin_memory=True): 包住前向。实践中常只 offload 一部分层，或与重算混用。图中时间为示意（拷贝取计算的 1.6 倍）。

交互

拷贝能不能藏在计算后面

一层激活（b=1）
一层前向耗时
D2H 拷贝耗时
拷贝 / 计算

示意估算：激活 \(34\,sbh\) 字节，前向 FLOPs \(24\,sbh^2 + 4\,bs^2h\)，按 989 TFLOPS 的 50% 算；拷贝按 50 GB/s 算。两者都和 \(b\) 成正比，比值只取决于 \(h\) 和 \(s\)：模型越宽，每字节激活背后的计算越多，越容易藏住。

流水线并行

1F1B：第一级要扛 p 个 micro-batch 的激活

4 级流水线、8 个 micro-batch 的 1F1B 调度。级编号从 0 开始，第 \(i\) 级在稳态时最多持有 \(p - i\) 个 micro-batch 的激活：第一级最多。

第一级有 \(L/p\) 层、持有 \(p\) 份，合起来正好是 \(L\) 层的激活，和 \(p\) 无关：流水线切得再细，第一级的激活也不减少。

加上重算：每份在途激活只剩层输入，变得很细；反向块里多出一段重算（橙色），整条流水线变长。也可以只对放不下的那几个 micro-batch 重算。

FSDP

重算发生在参数已经 gather 好的时候

FSDP 每层的参数平时是分片的：前向前 all-gather，用完释放；反向前再 all-gather 一次，算完梯度 reduce-scatter。

在图中的 FSDP 调度里，把 checkpoint 边界对齐分片单元，让重算复用反向已 gather 的参数，就能避免额外 all-gather。具体包装顺序和 reshard 策略会影响参数生命周期，用通信时间线核对。

FSDP 切的是参数、梯度和优化器状态，不切激活：每张卡的激活只取决于它自己的 micro-batch。能压激活的是重算、offload、序列 / 上下文并行。

PyTorch

torch.utils.checkpoint

from torch.utils.checkpoint import checkpoint

class Model(nn.Module):
    def forward(self, x):
        for blk in self.blocks:
            if self.training:
                # 只存 blk 的输入，反向前重跑 blk.forward
                x = checkpoint(blk, x, use_reentrant=False)
            else:
                x = blk(x)
        return self.head(self.norm(x))

# 或者按段：每段若干 block，只存段边界
from torch.utils.checkpoint import checkpoint_sequential
y = checkpoint_sequential(nn.Sequential(*blocks), 4, x,
                          use_reentrant=False)

右边是 autograd 保存的张量：普通前向每个 block 留下一堆；包上 checkpoint 后主要保留输入。checkpoint_sequential 的最后一段仍按普通前向执行。

务必显式写 use_reentrant=False（新版本不写会报警告）。非 reentrant 版本支持 torch.autograd.grad、不要求输入 requires_grad、能处理嵌套结构里的张量，重算到所需激活就提前停下。

dropout 的随机数默认会被保存并在重算时恢复（preserve_rng_state=True），两次前向结果一致。

语义 · non-reentrant 与重算范围

函数可能执行不止一次，而且不一定跑到底

显式使用 use_reentrant=False

non-reentrant 会记录 autograd 图，支持更灵活的反向调用；重算在所需张量恢复后可以提前停止。不要靠函数尾部执行次数维持状态。

把状态变更放在边界外

更新全局计数器、修改 cache、消耗外部迭代器或原地改输入，都可能使重算走到不同数值或控制分支。让被包函数只由显式输入决定。

默认 determinism_check 比较重算张量的 shape、dtype、device 等元信息，不会替你证明所有数值和梯度相同。真实验证要对比 loss、输出与参数梯度。

随机性 · dropout 不能重新抽一次

同一次前向的反向，需要同一张掩码

dropout(x) = x ⊙ mask / (1−p)
dL/dx = dL/dy ⊙ mask / (1−p)

取 p=½，x=[2, 4]，原 mask=[1, 0]，输出 [4, 0]。若 dy=[1, 1]，正确 dx=[2, 0]；重算若误抽 [0, 1]，就会得到 [0, 2]。

保存 RNG 状态

checkpoint 默认会保存并恢复相关随机状态，使重算与原调用的随机操作一致。可重复训练与不重复消耗外部随机流是两个相关目标。

仍要管住外部随机源

Python random、NumPy RNG、自定义 generator、函数里迁移到新设备等，不能想当然地视为都已覆盖；把随机结果显式传入最清楚。

实现 · 带 dropout 的对照实验

比较同一份参数、输入与随机种子

plain = copy.deepcopy(model)
checked = copy.deepcopy(model)
# 每次 run：zero_grad，克隆相同输入，重置相同 RNG seed
y0 = plain(x0)
y1 = checkpoint(checked, x1, use_reentrant=False,
                preserve_rng_state=True)
y0.square().mean().backward()
y1.square().mean().backward()
torch.testing.assert_close(y0, y1)
for p0, p1 in zip(plain.parameters(), checked.parameters()):
    torch.testing.assert_close(p0.grad, p1.grad)

完整脚本把两次 run 分开并分别重置随机种子；这里展示需要比较的对象。CPU 即可运行，安装 PyTorch 后执行 python3 torch_example.py。

下载完整 PyTorch 对照实验
按算子选择

选择性 checkpoint：存矩阵乘，重算逐元素算子

import functools, torch
from torch.utils.checkpoint import (checkpoint,
    create_selective_checkpoint_contexts, CheckpointPolicy)

ops_to_save = {
    torch.ops.aten.mm.default,
    torch.ops.aten._scaled_dot_product_flash_attention.default,
}

def policy(ctx, op, *args, **kwargs):
    if op in ops_to_save:
        return CheckpointPolicy.MUST_SAVE      # 贵：存
    return CheckpointPolicy.PREFER_RECOMPUTE   # 便宜：重算

ctx_fn = functools.partial(create_selective_checkpoint_contexts, policy)
y = checkpoint(blk, x, use_reentrant=False, context_fn=ctx_fn)

整层重算是"全有或全无"。更细的做法是逐个算子决定：policy 函数拿到每个 aten 算子，返回存还是重算。

矩阵乘和 attention 计算量大，结果存下；norm、激活函数、逐元素乘这些便宜的算子在反向前重算。反向时只需重跑便宜的那部分。

全景

各种做法在显存 / 计算平面上的位置

横轴是额外计算，纵轴是剩下的激活显存（以不重算为 100%；按 GPT-3 175B 形状、标准 attention 估算）。左下角最好。

先拿便宜的：在这个形状下，FlashAttention / attention core 重算用少量额外计算省掉约七成激活；剩下的再用部分层重算、逐层重算或 offload 补，按显存缺口决定重算多少层。

测量 · 同时观察峰值与训练吞吐

从相同训练状态开始计时

# warmup 若干完整 step，让编译和优化器状态分配完成
torch.cuda.synchronize()
torch.cuda.reset_peak_memory_stats()
start = time.perf_counter()
train_step()  # 含 zero_grad、forward、backward、optimizer.step
torch.cuda.synchronize()
seconds = time.perf_counter() - start
allocated = torch.cuda.max_memory_allocated()
reserved = torch.cuda.max_memory_reserved()

每种策略在相同模型、micro-batch、序列长与精度下独立测量；最好用单独进程保持初始化一致。allocated 是活跃张量口径，reserved 还含分配器缓存。

分布式训练还应同步计时边界，观察最慢 rank；报告 tokens/s，并确认 loss 与梯度保持预期。

排查 · OOM 消失后，训练还要正常

把显存、梯度、吞吐分别验一遍

现象优先检查有用的对照
checkpoint 后输出相同，梯度不同dropout RNG、状态变更、原地写固定种子比较每层参数梯度
前向省了，反向仍 OOM重算段太长、workspace / all-gather 峰值缩短段长并查看峰值时间线
reserved 不降，allocated 降了分配器保留已申请内存比较 peak allocated 与 memory snapshot
开启后速度损失远超估计通信重叠、反复 gather、编译分段对照计算/通信 profiler 时间线
两种策略 loss 不可比训练状态或样本数不同从相同 checkpoint、同一批数据开始

“没有报错”只说明执行完成；工程验收至少包含数值一致性、内存目标与吞吐目标。

练习 1 · 独立分析

最优段长为什么不是 √48？

48 层，入口 c=16 MiB，每层激活 A=128 MiB。用 ceil(48/k)c+kA 计算 k=1、2、3、4，并选择候选段长。
展开参考解答

① k=1：48×16+128 = 896 MiB；k=2：24×16+256 = 640 MiB。

② k=3：16×16+384 = 640 MiB；k=4：12×16+512 = 704 MiB。

③ 连续解 √(48×16/128)=√6≈2.45；2、3 都应实测。√48 只对应 c=A 的特例，不能用于这份账本。

先独立计算，再展开解答。

练习 2 · 独立分析

省显存之后，什么时候值得扩大 batch？

教学测量：原 micro-batch=1，一个 step 用 80 ms；checkpoint 后可用 micro-batch=2，一个 step 用 130 ms。序列长相同、无梯度累积。样本吞吐提高多少？比较学习效果还需固定什么？
展开参考解答

① 原吞吐 1/0.080 = 12.5 样本/s；新吞吐 2/0.130 ≈ 15.385 样本/s，提高约 23.1%。

② 单 step 更慢，不等于单位样本训练更慢。还要观察真实 GPU 利用率与通信开销。

③ 若要比较同一训练配方，保持有效全局 batch 可比：调整梯度累积，核对 loss 归一化、学习率策略与 optimizer update 次数。

先独立计算，再展开解答。

练习 3 · 独立分析

一次重算为什么改变了梯度？

函数含 dropout(p=0.5) 和 counter += 1；原掩码 [1,0]，重算掩码 [0,1]，上游 dy=[1,1]。求错误 dx，并说明两类修复。
展开参考解答

① 错误 dx = [0,2]；原前向对应的正确 dx 是 [2,0]。输出相同的对照不能替代梯度对照。

② 对 checkpoint 管理的随机源保留 RNG 状态；其他随机源用显式掩码或自行可重复的状态管理。

③ 把 counter 等外部副作用移出被重算函数；non-reentrant 可能提前停止，也不能依赖“正好执行两次”。

先独立计算，再展开解答。

动手 · 手写一条可重算的反向链

保存版与分段版，梯度必须对得上

python3 lab.py --layers 12 --segment 3
# 再试：--segment 1、--segment 4、--segment 12

NumPy 实现 hᵢ₊₁=tanh(hᵢWᵢ)。分段版只保留入口，反向恢复当前段；与完整保存版比较输出/梯度，再用中心差分独立验证每层一个参数。

默认：重算 12 层；激活峰值 13 → 8 份
差分误差 < 1e−7，最终打印 PASS

份数只计激活数组，含保留输出；不含参数、梯度和 matmul 工作区。修改题：预测 k 太大和太小时峰值分别由谁主导。

下载 lab.py · NumPy / CPU
小结

小结

  • 激活是前向产生、反向要用的张量；图中的保存量先堆积再释放。实际总显存还包括权重、梯度与临时工作区。
  • 标准 attention 的账本：每层 \(sbh(34 + 5as/h)\) 字节；GPT-3 175B 形状下 96 层约 275 GB。
  • checkpoint：只存层输入，反向前重算；逐层重算后峰值 ≈ \(L\) 个输入 + 1 层，计算 +33%。
  • 分段的最优段长 \(k^*=\sqrt{nc/A}\)；checkpoint 和一层一样大时就是 √n，Transformer 里接近每层一个。
  • 选择性重算 attention core：省 65%–70%，只多 1.6%–2.7% FLOPs；FlashAttention 在 kernel 内做了同样的事。
  • offload 不花算力但受 PCIe 带宽限制；模型越宽越容易藏住。
  • 1F1B 的第一级要存 \(L\) 层的激活；FSDP 不切激活，重算要和 FSDP 单元对齐；PyTorch 里用 use_reentrant=False，需要时再按算子选。
参考

参考与延伸阅读

系列 · 动画图解 LLM 系统

继续阅读

训练省显存的另一半是优化器状态和并行切分，见系列里的相邻篇目

← 回到博客