做深度学习训练的人大概都经历过这种时刻模型跑到一半显存爆了或者机器被抢占、进程被意外杀掉几个小时的训练成果瞬间归零。尤其是现在动辄几十亿参数的大模型单卡根本放不下梯度检查点Gradient Checkpointing几乎成了标配技能。PyTorch 里的torch.utils.checkpoint这个模块名字看着不起眼但它是用时间换空间这套思路里最实用的一把工具。这篇内容我想把 Checkpoint 机制从原理到实操完整拆一遍包括它到底省了什么、代价在哪、什么场景该用、什么场景用了反而拖慢训练以及我在实际项目里踩过的那些坑。不管你是刚接触显存优化还是已经在多卡训练里被 OOM 折磨过几轮应该都能从里面找到能直接抄作业的东西。1. 先搞清楚 Checkpoint 到底在省什么1.1 反向传播的显存账本要理解 Checkpoint得先把普通反向传播的显存开销算清楚。前向传播时每一层的中间激活值activation都要保留下来因为反向传播计算梯度时需要用到它们。假设一个网络有 N 层每层的激活值大小是 A那么光是存激活就要 N×A 的显存。层数越深、batch 越大、序列越长这个数字涨得越快。举个具体的例子。一个 Transformer 层hidden size 设为 1024序列长度 2048batch 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_sequentialPyTorch 这个模块主要暴露两个入口。checkpoint(function, *args)是最基础的用法你把一个函数通常是一个 module 的 forward 或者一段计算逻辑传进去它会在前向时用torch.no_grad()跑一遍不存中间激活反向时再重新计算。checkpoint_sequential(modules, segments, input)则是针对nn.Sequential的便捷封装把一串 module 按 segments 分成若干段每段作为一个检查点单元。先看基础用法。假设你有一个自定义的 blockimport 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_firstTrue) 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_gradTrue).cuda() # 普通前向 out_normal block(x) # checkpoint 前向 out_ckpt checkpoint(block, x, use_reentrantFalse)注意这里传的是block这个 module 本身checkpoint 内部会调用它的__call__。use_reentrantFalse这个参数很关键后面单独讲。checkpoint_sequential的用法则是from torch.utils.checkpoint import checkpoint_sequential model nn.Sequential(*[HeavyBlock(1024) for _ in range(12)]).cuda() out checkpoint_sequential(model, segments4, inputx, use_reentrantFalse)segments4表示把这 12 个 block 分成 4 段每段 3 个 block 作为一个检查点单元。段数越多显存省得越多但重算次数也越多。这个参数没有万能值得根据显存余量和速度要求调。2.2 use_reentrant 参数为什么必须关注use_reentrant是 PyTorch 1.11 之后引入的参数默认值在旧版本是 True新版本逐渐引导用户用 False。这两个模式的区别不是细节而是会影响你的代码能不能正常跑。reentrant 模式True是早期实现它依赖torch.autograd.Function的机制对输入有比较严格的要求传给 checkpoint 的输入必须至少有一个requires_gradTrue否则反向时不会触发重算梯度直接断掉。而且它不支持一些高级特性比如在 checkpoint 内部调用.detach()或者某些原地操作会出问题。非 reentrant 模式False是重新设计的实现基于torch.autograd.graph.saved_tensors_hooks兼容性更好支持更多场景包括嵌套 checkpoint、在 checkpoint 内部使用torch.autograd.grad等。官方现在推荐一律用use_reentrantFalse。我踩过的一个坑早期代码里没写这个参数升级 PyTorch 后行为变了某些分支下梯度莫名其妙变成 None。排查了半天才发现是 reentrant 模式对输入 requires_grad 的要求导致的。所以现在我的习惯是只要用 checkpoint显式写上use_reentrantFalse不给默认值留隐患。提示如果你的 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_layers24, dim1024): super().__init__() self.layers nn.ModuleList([HeavyBlock(dim) for _ in range(num_layers)]) self.num_layers num_layers def forward(self, x, ckpt_start8, ckpt_seg8): # 前段正常前向 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_reentrantFalse) x checkpoint(run_segment(ckpt_start ckpt_seg, self.num_layers), x, use_reentrantFalse) 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_typecuda, dtypetorch.bfloat16): logits model(input_ids, ckpt_start8, ckpt_seg8) 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 层 Transformerdim 1024seq 2048bf16 混合精度。配置batch size峰值显存单步耗时吞吐样本/秒无 checkpoint838.2G0.42s19.0无 checkpoint16OOM--全层 checkpoint816.5G0.61s13.1全层 checkpoint1631.8G1.18s13.6分段 checkpoint后16层819.1G0.55s14.5分段 checkpoint后16层1636.4G1.06s15.1从表里能看出几个规律。全层 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 常见问题速查表现象可能原因解决方向梯度为 Nonereentrant 模式输入无 requires_grad改用 use_reentrantFalseloss 不下降checkpoint 内有 detach/no_grad移除断流操作结果和不开不一致输入被原地修改改为非原地操作显存没降反升检查点划分过碎减少 segments 数量BN 统计量异常重算导致 running stats 更新两次排除 BN 或换 LayerNorm速度掉太多全层 checkpoint改为分段只包重层嵌套 checkpoint 报错reentrant 模式不支持嵌套改用 use_reentrantFalse这张表基本覆盖了我遇到过的绝大多数问题。实际排查时先看现象对号入座再按解决方向逐个试效率会高很多。4.4 和分布式训练的配合在多卡或分布式训练里用 checkpoint有几个额外注意点。DDPDistributedDataParallel下checkpoint 的重算发生在 backward 阶段而 DDP 的梯度同步也是在 backward 阶段触发的两者交错执行一般不会有冲突但要注意重算的计算量会计入 backward 时间影响梯度同步的时机。实测下来没有正确性问题只是通信和计算的重叠效率可能受一点影响。FSDPFullyShardedDataParallel下情况更复杂一些因为参数本身是分片存储的checkpoint 重算时需要把分片参数重新聚合。PyTorch 对 FSDP checkpoint 有专门的支持但要求用use_reentrantFalse并且建议用 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_fnlambda m: checkpoint_wrapper(m, checkpoint_implCheckpointImpl.NO_REENTRANT), check_fnlambda 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_memoryTrue) as prof: model(x) print(prof.key_averages().table(sort_bycuda_memory_usage, row_limit20))这段代码能列出各操作的显存占用按大小排序帮你定位哪些层是显存大户。定位之后针对性地包 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 方案跑稳了再单独写一篇聊聊。