做深度学习训练的人，几乎都遇到过那个经典的红色报错：CUDA out of memory。模型明明不算大，但batch size稍微调大一点，显存就爆了；想用单张卡跑个稍大一点的模型，参数还没装满，激活值先溢出来了。这时候很多人的第一反应是换显卡、加显存、或者砍batch size，其实在动手花钱之前，还有一个几乎每个框架都自带、但很多人没真正用明白的工具——梯度检查点（Gradient Checkpointing）。
梯度检查点的核心思想特别简单：用计算换显存。训练过程里，前向传播产生的中间激活值不再全部保存在显存里，而是只挑一部分关键节点存下来，反向传播需要用到的其他激活值，临时再重新算一遍。付出的代价是训练时间会变长，换来的是同样一块显卡能塞下更大的模型、更大的batch size或者更长的序列。
这篇文章适合正在被显存瓶颈卡住的人——不管是训练大模型、微调大模型，还是做长序列任务，只要你的GPU显存吃紧，都可以参考。我会把原理、实操、调优、踩坑这几个部分都聊一遍，尽量做到看完就能用。
1. 为什么训练时会爆显存：梯度检查点要解决的核心矛盾
1.1 显存到底被谁吃掉了
训练一个模型的显存占用，大概可以分成四块：模型参数、优化器状态、中间激活值、临时缓冲区。
很多人第一反应是"模型参数占大头"，但这个直觉在训练大模型时往往不太对。以Transformer为例，每个Transformer块的前向过程里，attention的Q、K、V、注意力矩阵、FFN的中间层输出、每层LayerNorm的统计量等等，都要在显存里留着，方便反向传播求梯度。序列越长、batch越大、层数越多，这一块的占用就越大。实测下来，很多场景里中间激活值的占用远超模型参数本身。
举个例子，一个参数量在10亿量级的模型，fp16参数大约2GB。可如果你的训练batch里有8条、每条2048个token，每个token的中间特征维度是1024，那么每层激活值轻松上百MB，几十层叠下来就是好几个GB。所以爆显存的时候，先别急着怪参数多，很多时候真正的"内存刺客"是那一大堆中间激活值。
1.2 反向传播为什么非要盯着激活值
训练的核心是反向传播，而反向传播依赖链式法则：计算某一层权重的梯度，需要用到该层前向传播时的输入和中间输出。比如Attention层要算对Q、K、V的梯度，就需要softmax之后的注意力分布、V矩阵这些中间结果。如果不存这些中间结果，梯度就是无源之水，根本没法算。
所以传统训练框架默认选择"空间换时间"：前向做完后，把每一层的中间激活值全部保留在显存里，反向传播的时候随取随用。对小模型这不是问题，但对大模型来说，空间代价越来越难承担。于是设计师就开始想：能不能倒过来，用时间换空间？
这就是梯度检查点的出发点了。
1.3 思路转变：把"存档"从每页改成每章
梯度检查点的做法很简单：前向传播时，不保留所有中间激活值，而是把网络切成若干段，每一段只保留输入端的张量，这个张量就叫"检查点"。反向传播算到某一层时，如果发现缺少激活值，就从对应的检查点出发，重新做一次这段的前向计算，把需要的中间结果现算出来。
这个概念可以类比成两种复习方式：一种是每读一页书就把这页的重点抄进笔记本，随时翻阅，显存占用高但省时间；另一种是只在每章开头折一个角，真需要哪页内容时，再回头翻一遍书。显存占用低了，但每次查资料都要多花点翻书时间。梯度检查点就是后者。
2. 梯度检查点的技术原理与算法细节
2.1 检查点到底存什么、丢什么
实际实现中，检查点的粒度可以有很多种。最省心的方案就是按层划分：每隔固定层数设一个检查点，检查点之间的中间激活值全部丢掉。
比如一个12层的Transformer网络，每隔3层设一个检查点。前向传播时，显存里只保留第0层输入、第3层输入、第6层输入、第9层输入这4个张量，其他中间激活值用完即弃。反向传播到第4到第6层时，需要这一段的中间输出，就从第3层的检查点张量出发，重新执行第4层到第6层的前向计算，得到需要的激活值后再继续反向。
这里有个关键点：重算是分段重算，不是从头到尾把整个网络再跑一遍。每次只重算当前检查点之后、当前需要的那一小段，时间成本可控，显存收益却很直接。
2.2 多出来的时间开销从哪里来
这个策略的本质，是让部分前向计算多跑一遍。如果每个检查点区间长度为L层，那么反向传播时每个区间大约要多算一次前向。整体来看，额外计算量大约相当于多跑了一次完整的前向传播。
有多慢？社区实测里，梯度检查点通常会带来20%到40%的训练耗时增加；如果检查点粒度设得太细，耗时甚至可能翻倍。但这件事得换个角度算账：显存占用可能从"完全跑不了"变成"勉强能跑"，这是0到1的区别，多花点时间去换训练的可能性，非常划算。
2.3 检查点间隔怎么选：不是越多越好
理论上检查点越多，显存占用越低，但计算开销越大。最佳间隔取决于模型结构、激活值大小、显存上限。实操中比较稳妥的选择是：
- Transformer模型每1到2个块放一个检查点，显存收益很明显，长序列场景尤其如此。
- 如果激活值特别大（比如长序列、大图像输入），可以每1个块放一个，激进省显存。
- 如果显存压力不大，可以放宽到每4个块放一个，优先保证训练速度。
我习惯先用默认配置跑一轮，观察两个指标：显存峰值和每步耗时。显存还有余量就放宽间隔，显存仍然紧张就收紧间隔，直到找到一个"显存刚好够、速度还能接受"的点。
另外还有更进阶的策略，比如选择性检查点：把重计算代价高的层的激活值保留下来，重算那些计算代价低的层。工程上大多数情况用框架的默认封装就够了，手动控制粒度主要是在两极场景里用——要么追求极致省显存，要么追求最高训练吞吐。
3. 实操：如何在自己的训练代码里开启梯度检查点
3.1 框架内置的调用方式
大多数情况你不需要从零实现，主流深度学习框架都内置了现成接口。以PyTorch为例，最核心的API是torch.utils.checkpoint.checkpoint。最简单的用法：把某个模块的前向函数传进去。
from torch.utils.checkpoint import checkpoint class MyBlock(nn.Module): def forward(self, x): # 这里是正常的模块前向逻辑 return x class MyModel(nn.Module): def __init__(self): super().__init__() self.block = MyBlock() def forward(self, x): # 用checkpoint包裹，前向时只保存输入，不保存中间激活 return checkpoint(self.block, x)还有更省事的checkpoint_sequential，可以直接按顺序对一组模块做分段检查点：
from torch.utils.checkpoint import checkpoint_sequential # 假设 layers 是一个包含若干子模块的 ModuleList output = checkpoint_sequential(layers, segments=len(layers) // 3, input=x)在很多预训练模型代码库里，通常会有一个更简单的开关，比如model.gradient_checkpointing_enable()或者类似名字的函数。打开这个开关后，库内部会在每个Transformer Block外面自动套上checkpoint逻辑。所以大多数情况下，你连前向代码都不用改，调个开关就行。
3.2 一个完整的开启流程示例
我从头到尾给你梳理一遍推荐流程：
- 先看当前显存占用构成。可以在训练脚本里临时存一个检查点文件，输出一下
torch.cuda.max_memory_allocated()，再对比模型参数量，确认瓶颈到底是不是激活值。 - 打开框架提供的梯度检查点开关，如果没有，就用
checkpoint手动包住重模块。 - 把batch size调大，或者换成更大的模型。这一步是目的，不是顺便。
- 观察训练速度。每步耗时增加在30%以内，一般可以接受；超过太多就调大检查点间隔。
- 做一次完整的短训（几十步），确认梯度数值正常、loss能下降，再启动长训。
开启后有一点要注意：不要在checkpoint包裹的前向函数里对输入张量做原地操作。因为重算时需要的是"原始输入"，如果你在第一次前向时把它改了，重算就会拿到一个被污染的张量，梯度直接算错。这种错误特别隐秘，loss可能会疯涨或者完全变成NaN。
3.3 分布式训练和常见加速库里的配合
数据并行场景下梯度检查点可以放心开。省下来的显存可以用于增大每卡batch size，每卡batch大了，梯度同步效率也更高。
最常用的组合是：优化器状态分片（比如ZeRO一类的方案） + 梯度检查点 + 自动混合精度。这三者是互补的：优化器状态分片解决的是参数和优化器状态占显存的问题，梯度检查点解决激活值占显存的问题，混合精度把整体精度降到fp16降低占用。结合起来，单卡能跑的模型上限会大幅提升。
这里有个容易误解的点：梯度检查点和推理的关系。推理阶段通常不需要开——因为推理不做反向传播，根本不保存中间激活值。如果推理阶段显存爆了，那问题一般出在批量处理太大或者KV Cache之类的缓存上，跟梯度检查点不是同一个赛道。
4. 常见问题与排查技巧实录
4.1 开了检查点，显存却没降多少
这个问题我见过不止一次。如果你开了检查点，显存占用纹丝不动，先检查这几件事：
第一，显存大头可能根本不在激活值。如果你用的是Adam优化器，优化器状态里要存两份fp32副本，显存占用轻松达到参数量的8倍以上。如果模型本身很大，优化器状态可能才是吃掉显存的真凶。这种情况下应该上混合精度或优化器状态分片，而不是死磕梯度检查点。
第二，确认开关真的生效了。有些框架的检查点开关是给"训练入口"用的，如果你改完配置后从某个中间状态继续训练，配置可能没有同步到当前运行的模型上。建议打印一下模型结构，看看目标模块外面是不是真的套上了checkpoint包装。
第三，用显存分析工具看看到底哪里在涨。nvidia-smi显示的进程占用是包含缓存在内的，实际要关注的是PyTorch缓存分配器的峰值指标，也就是torch.cuda.max_memory_allocated()，这才是有意义的参考数字。
4.2 训练速度慢了很多，正常吗
先给个正常范围：慢20%到40%都是正常的，毕竟前向多算了一遍。但如果慢了50%以上甚至翻倍，就要排查了。
最常见的原因是检查点粒度太密。每层都设检查点，意味着每层都要额外重算一次，耗时自然起飞。建议先退回"每隔几个块设一个检查点"的粗粒度配置。
另一种可能是你把检查点包在了非关键模块上。比如某个模块只有很小的中间激活，重算它的收益极低，但计算开销不小。这时应该把检查点挪到Attention、FFN这类"激活值大、层数深"的模块上。
还有一个容易忽略的点：如果你的网络里有一些自定义算子，重算时没法利用到某些现成的缓存结果，可能会触发额外的kernel执行，导致速度骤降。这种情况没有通用解法，只能逐个模块测试重算耗时。
4.3 和自动混合精度一起用，注意数值稳定性
梯度检查点和自动混合精度（AMP）是可以共存的，但有一个坑：重算部分的数值精度。
正常前向在AMP模式下会走fp16，省显存又加速。但某些对精度敏感的任务里，重算时如果仍然用fp16，经过一串算子以后误差会被放大，导致梯度抖动，训练不稳定。我自己就遇到过重算后loss正常下降、但某些层梯度偶发NaN的情况，排查到最后就是精度问题。
解决办法是让重算部分保持fp32，只有正常前向走AMP。多数框架的内置开关会处理这个细节，但如果你是自己用checkpoint手动包的，就要注意上下文管理，确保重算过程中dtype不会被AMP切换到fp16。
另外，AMP里的梯度缩放器（GradScaler）不需要额外改动，正常用就行。
4.4 随机性、dropout与"不可重入"模块的问题
这个坑比较隐蔽，但一旦踩中就特别容易让人崩溃：梯度检查点重算前向时，dropout会再执行一次，导致两次随机采样结果不一致，梯度计算就出错了。
框架内置实现一般会自动保存和恢复随机数生成器的状态，确保重算时产生的dropout mask和第一次一致。但如果你自己写了手动的重算逻辑，或者用了不兼容的第三方模块，这个保护可能失效。
还有一种情况更麻烦：模块里有非确定性算子。比如某些自定义CUDA kernel，每次执行结果可能有微小差异，这种模块就不适合做重算。遇到这种情况，只能把这类模块排除在检查点之外，保留它的激活值。
4.5 常见问题速查表
我把高频问题整理成一张表，方便你排查时快速对照。
| 现象 | 可能原因 | 解决方向 |
|---|---|---|
| 显存没降 | 激活值不是显存瓶颈 | 改用混合精度、优化器状态分片 |
| 显存没降 | 检查点开关未生效 | 检查模型包装状态、配置同步 |
| 速度慢到翻倍 | 检查点粒度过细 | 放宽到每4个块设一个检查点 |
| 梯度异常/NaN | 重算精度问题 | 重算部分保持fp32 |
| 梯度对不上 | dropout随机mask不一致 | 用框架内置检查点实现，锁定RNG状态 |
| 自定义模块报错 | 算子不可重入 | 排除在检查点之外 |
5. 与其他显存优化手段的对比与组合
5.1 不同手段的定位完全不同
显存优化的工具其实不少，但很多人把它们混为一谈。我列个表说清楚：
| 手段 | 省的是什么显存 | 速度影响 | 典型场景 |
|---|---|---|---|
| 梯度检查点 | 中间激活值 | 慢20%~40% | 长序列、大batch、深层模型 |
| 混合精度（fp16/bf16） | 参数、激活、梯度整体减半 | 通常变快 | 几乎适合所有训练 |
| 梯度累积 | 不省单步显存，等价于增大batch | 基本无影响 | 模拟大batch但显存不够时 |
| 优化器状态分片 | 优化器状态 | 有通信开销 | 多卡训练超大模型 |
| 参数/优化器状态卸载 | 优化器状态和参数 | 明显变慢 | 单卡硬跑超大模型 |
从这张表能看出，梯度检查点解决的是"激活值"这个特定瓶颈，其他手段各有各的适用面。排错时先定位是哪块显存爆了，再选对应工具，这才是效率最高的路径。
5.2 一个实用的组合方案示例
我个人的使用习惯是：混合精度 + 梯度检查点 + 梯度累积，三者一起上。这个组合几乎成了我训练大模型的默认起手式。
举个实际例子：假设你单卡24GB显存，要微调一个7B量级的模型。纯fp32训练一定跑不动，先开混合精度，显存立刻降一截；再开梯度检查点，激活值大幅缩减，batch size可以往上抬；如果还想继续加大有效batch size，就用梯度累积，每几步累积一次梯度再更新权重。
这样组合下来，原来只能塞下很小batch的模型，能在一个相对合理的batch size下跑起来了。代价是训练时间变长，但总比换显卡、等资源要实际得多。
分布式多卡场景里，再叠加优化器状态分片，单卡压力进一步降低。需要注意的是，优化器状态分片本身有通信开销，所以小规模多卡里收益不一定大；卡数足够多、模型足够大时，收益才明显。
5.3 什么时候不推荐梯度检查点
不是所有场景都适合开梯度检查点。如果你的显存还有不少余量，开它只会白白增加训练耗时。另外，如果模型里重计算成本极高，或者存在无法安全重算的自定义算子，开检查点的性价比就很低。
还有一种情况：你的训练本身已经是计算密集型的，GPU算力跑满了，再开检查点只会让每步时间明显拉长。这时更合理的思路是降低激活值占用，比如缩短序列长度、用小batch配合梯度累积，或者优化注意力算子，而不是无脑开检查点。
最后再分享一点个人经验
玩了这么久训练调优，我最大的感受是：遇到显存问题，先诊断再动手，比盲目开优化手段有效得多。显存分析工具和max_memory_allocated这类指标，比凭感觉猜瓶颈靠谱一万倍。
另外，如果你开了梯度检查点后发现训练变慢得厉害，有个小技巧：不要整块网络统一设检查点，试着只包住Attention和FFN这类"激活值特别大"的子模块，往往能在显存收益和训练速度之间找到更好的平衡点。我拿这个思路处理过一次长序列训练，显存降了快一半，速度只慢了十几个百分点，效果比粗粒度全量开启好很多。
梯度检查点不是一个万能开关，但它绝对是显存工具箱里最有性价比的一件工具。希望这篇文章能帮你少走点弯路，把省下来的显存用在真正有用的地方。
