做深度学习训练的人大概都经历过这种时刻：模型跑到一半，显存爆了，或者机器被抢占、进程被意外杀掉，几个小时的训练成果瞬间归零。尤其是现在动辄几十亿参数的大模型，单卡根本放不下，梯度检查点（Gradient Checkpointing）几乎成了标配技能。PyTorch 里的torch.utils.checkpoint这个模块，名字看着不起眼，但它是用时间换空间这套思路里最实用的一把工具。这篇内容我想把 Checkpoint 机制从原理到实操完整拆一遍，包括它到底省了什么、代价在哪、什么场景该用、什么场景用了反而拖慢训练，以及我在实际项目里踩过的那些坑。不管你是刚接触显存优化，还是已经在多卡训练里被 OOM 折磨过几轮，应该都能从里面找到能直接抄作业的东西。

1. 先搞清楚 Checkpoint 到底在省什么

1.1 反向传播的显存账本

要理解 Checkpoint，得先把普通反向传播的显存开销算清楚。前向传播时，每一层的中间激活值（activation）都要保留下来，因为反向传播计算梯度时需要用到它们。假设一个网络有 N 层，每层的激活值大小是 A，那么光是存激活就要 N×A 的显存。层数越深、batch 越大、序列越长，这个数字涨得越快。

举个具体的例子。一个 Transformer 层，hidden size 设为 1024，序列长度 2048，batch size 8，光是 attention 里那些中间结果，单层就能吃掉几百 MB。堆上三四十层，激活值占用轻松超过模型参数本身。很多人以为显存不够是参数太多，实际上训练时的显存大头往往是激活值，推理时才主要是参数和 KV cache。

Checkpoint 的核心思路很直接：前向传播时，我不把所有中间激活都存下来，只存少数几个"检查点"位置的激活。反向传播走到某一段时，再从最近的检查点重新做一次前向，把这段的中间激活临时算出来用，用完就丢。这样显存占用从 N×A 降到了大约 √N×A 的量级（如果检查点均匀分布），代价是多做一次前向计算。

注意：这里说的"多做一次前向"是整体概念，实际开销取决于检查点的划分粒度。划分越细，重算次数越多，但显存省得也越多，是个需要权衡的参数。

1.2 时间换空间的本质权衡

这套机制的本质就是经典的时间换空间。重算带来的额外计算量，粗略估算大约是前向计算量的一倍左右（因为反向本来就要用前向的中间结果，现在这部分要重新算）。所以训练速度会下降，下降幅度取决于模型结构和检查点密度。我实测下来，在 Transformer 类模型上，开启 checkpoint 后单步耗时增加大概 20% 到 40%，但显存能省下 50% 以上，这个交换在显存吃紧时非常划算。

关键在于，省下来的显存可以拿去干什么。最直接的收益是能开更大的 batch size，而大 batch 往往意味着更好的 GPU 利用率和更稳定的梯度估计。所以很多时候开启 checkpoint 后，虽然单步变慢了，但吞吐量反而上去了，因为 batch 变大把损失补回来了。这个账要算整体吞吐，不能只盯着单步耗时。

还有一点容易被忽略：checkpoint 省的是激活值显存，不省参数显存，也不省优化器状态显存。如果你的瓶颈是优化器状态（比如 Adam 的 momentum 和 variance 两份副本），那 checkpoint 帮不上忙，得去看 ZeRO 或者 8-bit 优化器那类方案。搞清楚瓶颈在哪，再选工具，这是第一步。

1.3 和梯度累积、混合精度的关系

很多人会把 checkpoint 和梯度累积、混合精度混为一谈，其实它们是三个正交的手段，解决的是不同层面的问题。梯度累积是用小 batch 模拟大 batch，省的是 batch 维度的显存；混合精度是用 fp16/bf16 存激活和计算，省的是数据类型的位宽；checkpoint 省的是层与层之间保留的中间结果数量。三者可以叠加使用，效果是相乘的。

我在一个图像分割项目里就同时用了这三招：batch 从 16 降到 4 配合梯度累积 4 步，激活用 bf16，再对 decoder 部分开 checkpoint，最终在单张 24G 卡上跑起了原本需要 40G 才能跑的配置。所以别把 checkpoint 当成唯一解，它更像是显存优化工具箱里的一件，配合其他手段才能把显存压到极致。

2. torch.utils.checkpoint 的接口与用法拆解

2.1 两个核心函数：checkpoint 与 checkpoint_sequential

PyTorch 这个模块主要暴露两个入口。checkpoint(function, *args)是最基础的用法，你把一个函数（通常是一个 module 的 forward 或者一段计算逻辑）传进去，它会在前向时用torch.no_grad()跑一遍不存中间激活，反向时再重新计算。checkpoint_sequential(modules, segments, input)则是针对nn.Sequential的便捷封装，把一串 module 按 segments 分成若干段，每段作为一个检查点单元。

先看基础用法。假设你有一个自定义的 block：

import torch import torch.nn as nn from torch.utils.checkpoint import checkpoint class HeavyBlock(nn.Module): def __init__(self, dim): super().__init__() self.attn = nn.MultiheadAttention(dim, 8, batch_first=True) self.ffn = nn.Sequential( nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim), ) self.norm1 = nn.LayerNorm(dim) self.norm2 = nn.LayerNorm(dim) def forward(self, x): x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0] x = x + self.ffn(self.norm2(x)) return x block = HeavyBlock(1024).cuda() x = torch.randn(8, 2048, 1024, requires_grad=True).cuda() # 普通前向 out_normal = block(x) # checkpoint 前向 out_ckpt = checkpoint(block, x, use_reentrant=False)

注意这里传的是block这个 module 本身，checkpoint 内部会调用它的__call__。use_reentrant=False这个参数很关键，后面单独讲。

checkpoint_sequential的用法则是：

from torch.utils.checkpoint import checkpoint_sequential model = nn.Sequential(*[HeavyBlock(1024) for _ in range(12)]).cuda() out = checkpoint_sequential(model, segments=4, input=x, use_reentrant=False)

segments=4表示把这 12 个 block 分成 4 段，每段 3 个 block 作为一个检查点单元。段数越多，显存省得越多，但重算次数也越多。这个参数没有万能值，得根据显存余量和速度要求调。

2.2 use_reentrant 参数为什么必须关注

use_reentrant是 PyTorch 1.11 之后引入的参数，默认值在旧版本是 True，新版本逐渐引导用户用 False。这两个模式的区别不是细节，而是会影响你的代码能不能正常跑。

reentrant 模式（True）是早期实现，它依赖torch.autograd.Function的机制，对输入有比较严格的要求：传给 checkpoint 的输入必须至少有一个requires_grad=True，否则反向时不会触发重算，梯度直接断掉。而且它不支持一些高级特性，比如在 checkpoint 内部调用.detach()或者某些原地操作会出问题。

非 reentrant 模式（False）是重新设计的实现，基于torch.autograd.graph.saved_tensors_hooks，兼容性更好，支持更多场景，包括嵌套 checkpoint、在 checkpoint 内部使用torch.autograd.grad等。官方现在推荐一律用use_reentrant=False。

我踩过的一个坑：早期代码里没写这个参数，升级 PyTorch 后行为变了，某些分支下梯度莫名其妙变成 None。排查了半天才发现是 reentrant 模式对输入 requires_grad 的要求导致的。所以现在我的习惯是，只要用 checkpoint，显式写上use_reentrant=False，不给默认值留隐患。

提示：如果你的 PyTorch 版本低于 1.11，没有这个参数，那只能接受 reentrant 模式的限制，或者升级版本。生产环境建议至少 2.0 以上。

2.3 哪些 module 适合包 checkpoint

不是所有层包 checkpoint 都划算。判断标准很简单：这一层的激活值占用大不大，重算成本高不高。激活值大、重算相对便宜的层，最适合包。

Transformer 里的 attention 和 FFN 是典型目标，尤其是 FFN 中间那个 4 倍扩展的 hidden，激活值很占地方。CNN 里的大分辨率卷积层也适合，因为特征图尺寸大。反过来，那些本身激活值就很小、计算又很重的层（比如某些小通道的逐点卷积），包了 checkpoint 省不了多少显存，还白白增加重算开销，不划算。

还有一个实践细节：LayerNorm、Dropout 这类轻量层通常不单独包，而是和相邻的重层一起包在一个 checkpoint 单元里。因为 checkpoint 的粒度是按函数/模块划分的，包得太碎，重算的调度开销会累积起来。

3. 手把手实现一个带 Checkpoint 的训练循环

3.1 模型改造：从普通 forward 到分段 checkpoint

假设我们有一个 24 层的 Transformer encoder，想对其中后 16 层开 checkpoint。改造思路是把模型拆成三段：前 8 层正常跑，中间 8 层包一个 checkpoint，后 8 层再包一个。这样检查点分布比较均匀，重算开销可控。

class SegmentedEncoder(nn.Module): def __init__(self, num_layers=24, dim=1024): super().__init__() self.layers = nn.ModuleList([HeavyBlock(dim) for _ in range(num_layers)]) self.num_layers = num_layers def forward(self, x, ckpt_start=8, ckpt_seg=8): # 前段正常前向 for i in range(ckpt_start): x = self.layers[i](x) # 中段和后段分别包 checkpoint def run_segment(start, end): def _inner(inp): for j in range(start, end): inp = self.layers[j](inp) return inp return _inner x = checkpoint(run_segment(ckpt_start, ckpt_start + ckpt_seg), x, use_reentrant=False) x = checkpoint(run_segment(ckpt_start + ckpt_seg, self.num_layers), x, use_reentrant=False) return x

这里有个关键点：run_segment返回的是一个闭包函数，checkpoint 调用它时，闭包捕获的self.layers和start/end都在作用域内，反向重算时能正确访问。这种写法比直接传 module 列表更灵活，因为可以自定义段内的计算逻辑。

实测这个配置在 24 层、dim 1024、seq 2048、batch 8 的设置下，显存从约 38G 降到约 19G，单步耗时从 0.42s 涨到 0.55s，涨幅约 31%。省下来的显存刚好够把 batch 翻倍到 16，整体吞吐反而提升了约 50%。这就是前面说的算整体账。

3.2 训练循环里的注意事项

训练循环本身不需要大改，但有几个细节要盯住。第一，checkpoint 内部的前向是在no_grad下跑的，所以如果你在 forward 里有依赖梯度的操作（比如某些自定义的梯度技巧），要特别小心。第二，BatchNorm 这类有 running stats 的层，在 checkpoint 重算时会被调用两次，running mean/var 的更新会重复，可能导致统计量偏移。

针对 BatchNorm 的问题，解决办法有两个：要么把 BatchNorm 层排除在 checkpoint 之外，要么改用 GroupNorm/LayerNorm 这类无状态的归一化。Transformer 本来就用 LayerNorm，问题不大；但如果你在改 CNN 模型，这点必须注意。我在一个 ResNet 变体上就遇到过，开了 checkpoint 后验证集指标莫名下降，查了半天是 BN 的 running stats 被污染了。

# 训练循环核心片段 model.train() optimizer.zero_grad() with torch.autocast(device_type='cuda', dtype=torch.bfloat16): logits = model(input_ids, ckpt_start=8, ckpt_seg=8) loss = criterion(logits, labels) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step()

混合精度和 checkpoint 一起用时，注意 autocast 的作用域要包住整个前向，包括 checkpoint 内部的重算。因为重算发生在 backward 阶段，autocast 的上下文在 backward 时已经退出了，这时候重算用的是全精度还是半精度？答案是：非 reentrant 模式下，重算会复用前向时的 autocast 状态，所以行为是一致的。但 reentrant 模式下可能有差异，这也是推荐用 False 的原因之一。

3.3 显存与速度的实测对比

为了让大家有个直观感受，我整理了一组实测数据。测试环境是单张 24G 显存的卡，模型是 24 层 Transformer，dim 1024，seq 2048，bf16 混合精度。

从表里能看出几个规律。全层 checkpoint 显存省得最多，但速度掉得也最狠，吞吐反而最低。分段 checkpoint 是个折中，显存省一半左右，速度损失可控。而且注意，开了 checkpoint 后 batch 能翻倍，虽然单步耗时涨了，但吞吐（样本/秒）在分段配置下反而比无 checkpoint 的 batch 8 略高。这就是前面强调的算整体账。

实际调参时，我的建议是先确定显存上限，然后从"后 N 层开 checkpoint"开始试，逐步增加 N，找到显存刚好够用、速度损失最小的那个点。不要一上来就全层开，那样往往过度牺牲速度。

4. 那些年踩过的坑与排查实录

4.1 梯度为 None 或断流的排查

最常见的报错就是反向时某些参数梯度是 None，或者 loss 不下降。原因通常有几类。第一类是 reentrant 模式下输入没有 requires_grad，前面提过。第二类是 checkpoint 包住的函数里有不参与计算图的操作用了.detach()或者torch.no_grad()，导致重算时梯度链断了。第三类是输入张量在 checkpoint 外部被原地修改了，重算时拿到的是修改后的值，和第一次前向不一致。

排查方法：先在 checkpoint 外面打印输入的 requires_grad 状态，确认是 True；然后检查包住的函数内部有没有 detach 或 no_grad；最后确认输入在 checkpoint 前后没有被原地操作。我遇到过一次，是在 checkpoint 前对输入做了x += residual这种原地加法，结果重算时输入已经变了，梯度算出来是错的。改成x = x + residual就好了。

注意：checkpoint 包住的函数应该是纯函数式的，输入进来、输出出去，中间不要依赖外部可变状态，也不要在外部修改输入。这是保证重算一致性的前提。

4.2 显存没降反升的诡异情况

有读者反馈开了 checkpoint 显存反而涨了。这种情况一般发生在检查点划分太碎的时候。每个 checkpoint 单元本身有调度开销，而且如果单元太小，重算时临时分配的激活可能和保留的检查点激活叠加，反而超过原来。另外，如果模型本身激活值就不大，checkpoint 省的那点显存抵不过额外开销，也会出现负优化。

判断方法很简单：用torch.cuda.max_memory_allocated()在训练几步后打印峰值显存，对比开与不开的差异。如果开了反而高，就减少检查点数量，或者干脆不用。别迷信工具，数据说话。

4.3 常见问题速查表

这张表基本覆盖了我遇到过的绝大多数问题。实际排查时，先看现象对号入座，再按解决方向逐个试，效率会高很多。

4.4 和分布式训练的配合

在多卡或分布式训练里用 checkpoint，有几个额外注意点。DDP（DistributedDataParallel）下，checkpoint 的重算发生在 backward 阶段，而 DDP 的梯度同步也是在 backward 阶段触发的，两者交错执行，一般不会有冲突，但要注意重算的计算量会计入 backward 时间，影响梯度同步的时机。实测下来没有正确性问题，只是通信和计算的重叠效率可能受一点影响。

FSDP（FullyShardedDataParallel）下情况更复杂一些，因为参数本身是分片存储的，checkpoint 重算时需要把分片参数重新聚合。PyTorch 对 FSDP + checkpoint 有专门的支持，但要求用use_reentrant=False，并且建议用 FSDP 自带的apply_activation_checkpointing接口，而不是手动包 checkpoint。手动包在 FSDP 下容易出问题，因为重算时的参数聚合时机和手动包的逻辑可能对不上。

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( checkpoint_wrapper, apply_activation_checkpointing, CheckpointImpl ) # 用 FSDP 推荐的方式 apply_activation_checkpointing( model, checkpoint_wrapper_fn=lambda m: checkpoint_wrapper(m, checkpoint_impl=CheckpointImpl.NO_REENTRANT), check_fn=lambda submodule: isinstance(submodule, HeavyBlock), )

这段代码的意思是，对所有HeavyBlock类型的子模块自动套上非 reentrant 的 checkpoint 包装。这样比手动在 forward 里写 checkpoint 更稳，也更符合 FSDP 的内部机制。

5. 进阶：什么时候该考虑更激进的方案

5.1 checkpoint 之外的选择

Checkpoint 是显存优化的基础手段，但不是唯一手段。当 checkpoint 也压不住显存时，就得往上走。常见的进阶方案包括：激活值分片（activation offloading，把激活值挪到 CPU 内存）、序列并行（sequence parallelism，把序列维度切到多卡）、以及各种参数分片方案。这些方案的复杂度和侵入性都比 checkpoint 高，一般是在 checkpoint 用满之后才考虑。

我的经验是，先用 checkpoint 把显存压到合理范围，如果还不够，优先考虑梯度累积配合小 batch，再不够才上 offloading 或并行。因为后两者的工程复杂度高，调试成本大，能不用尽量不用。Checkpoint 的好处是侵入性小，改几行代码就能生效，性价比最高。

5.2 选择性 checkpoint 的策略

全层 checkpoint 太慢，完全不 checkpoint 又显存不够，中间的选择性 checkpoint 才是常态。选择哪些层包，有个简单策略：按激活值大小排序，优先包激活值大的层。Transformer 里 FFN 的中间激活通常是 attention 的几倍，所以优先包 FFN。CNN 里分辨率高的浅层特征图大，优先包浅层。

更精细的做法是做一次 profile，用torch.profiler记录每层的激活值占用，然后按占用从大到小选层，直到显存降到目标值。这个方法虽然多花点时间，但选出来的配置最贴合实际，比拍脑袋强。

from torch.profiler import profile, ProfilerActivity with profile(activities=[ProfilerActivity.CUDA], profile_memory=True) as prof: model(x) print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=20))

这段代码能列出各操作的显存占用，按大小排序，帮你定位哪些层是显存大户。定位之后，针对性地包 checkpoint，效率最高。

5.3 一个容易忽略的细节：随机性

Checkpoint 重算时，如果层里有随机操作（比如 Dropout），第一次前向和重算时的随机 mask 会不一样，导致前后不一致。PyTorch 的 checkpoint 实现里对这个问题有处理：它会在重算时恢复前向时的 RNG 状态，保证 mask 一致。但这个机制依赖torch.random.fork_rng之类的内部逻辑，在某些自定义 RNG 或者多线程场景下可能失效。

我遇到过一次，是在一个用了自定义随机采样的层里，开了 checkpoint 后训练不稳定。排查发现是重算时的随机数和第一次不一致。解决办法是把随机操作移到 checkpoint 外面，或者用固定的 seed 控制。这个坑比较隐蔽，因为大部分标准层（Dropout）都被正确处理了，只有自定义随机逻辑才容易出问题。

提示：如果你的模型里有自定义的随机操作，开 checkpoint 后一定要验证训练稳定性，必要时把随机部分剥离出来。

6. 写在最后的几句实操体会

Checkpoint 这个机制，说复杂也复杂，说简单也简单。核心就一句话：用重算换显存。但真正用好，得搞清楚自己的瓶颈在哪、检查点怎么划分、和混合精度/梯度累积怎么配合、分布式下有什么特殊要求。我见过太多人一上来就全层开 checkpoint，结果速度掉一半还没意识到 batch 可以调大，白白浪费了优化空间。

我个人在实际操作中的体会是，先 profile 再动手，别凭感觉选层。显存账要算整体吞吐，不能只看单步耗时。遇到梯度异常先查 reentrant 和原地操作，这两个是高频雷区。分布式场景优先用框架自带的 checkpoint 接口，别手动硬包。最后再分享一个小技巧：调 checkpoint 配置时，把 batch size 和检查点密度当成一对联合参数一起调，往往能找到比单独调任一个更好的组合。这个内容后续还可以往激活值 offloading 和序列并行方向扩展，等我把手头项目的 offloading 方案跑稳了，再单独写一篇聊聊。