想象你在参加一场限时数学竞赛,草稿纸只有薄薄两页,而题目需要十二步演算。如果你把每一步的中间结果都工工整整抄在草稿纸上,第二页还没写到一半纸就用完了。聪明的学霸会怎么做?只在第 1、5、9 步的关键节点记下一个结果,其余中间过程算完就翻篇;等后面需要某个中间结果时,回头从最近的关键节点重新快速推演一遍。AI 训练大模型时面临的正是同样的窘境——GPU 显存就是那张草稿纸,模型每一层前向传播产生的“中间激活值”就是那些中间结果。当模型深达几十上百层时,草稿纸根本不够写。梯度检查点(Gradient Checkpointing)就是那位学霸的做题策略:只保存少数“检查点”层的激活,反向传播时把丢掉的激活重新算回来,用一点额外计算,换取成倍的显存空间。
掌握梯度检查点,是 AI 从业者突破单卡显存天花板、训练更大模型、开更大 batch 的必备杠杆——不动硬件就能把“装不下”变成“放得下”。
可视化演示
展示梯度检查点进阶三维度:间隔优化(每隔k层保存检查点反向重算k层显存从O(n)降至O(√n)最优间隔k=√L)、重计算正确性(前向保存checkpoint输入反向重算恢复激活RNG状态保存确保Dropout等随机操作复现梯度数值一致)、协同优化(检查点+ZeRO/FSDP参数分片+CPU Offload卸载多技术叠加将显存降至原始15%支持千亿参数训练),三模式切换理解大模型训练显存优化工程方案。
本节配有交互动效(进入视口后自动加载,约 3KB)。
'+C1+'
显存危机
回到数学竞赛的比喻:你发现真正的问题不只是草稿纸太少,而是每一步的中间结果都太占地方。训练大模型时,显存里住着四位“房客”:模型参数、梯度、优化器状态,以及最贪心的——前向传播中每一层产生的中间激活值。前面三位的大小基本等于模型参数量乘以一个常数,而激活值的大小却和层数、batch size、序列长度直接挂钩:层越多、batch 越大,这房客占的地盘就成倍膨胀。
在 AI 训练中,这四位房客共同构成了训练时的显存需求。如果一张 GPU 卡只有 16GB 显存,而模型本身就“重”得让草稿纸所剩无几,那么最常见的结果就是 OOM(显存溢出)——训练直接中断,就像你在草稿纸上写不下最后一步一样。
当草稿纸被中间步骤占满,真正用来思考新问题的空白越来越少——这就是大模型训练最常撞上的那堵墙:单卡显存装不下。而梯度检查点的全部智慧,就是针对最贪心的“激活值房客”下手。
- 参数显存模型权重本身占用的显存,比如一个 70 亿参数的模型以 16 位浮点存储,约需 14GB。
- 梯度显存反向传播时每个参数都需要一个梯度副本,大小与参数显存相同。
- 优化器显存Adam 等优化器要为每个参数保存动量与二阶动量,通常是参数显存的 2 倍。
- 激活显存前向传播每一层输出的中间张量都会留在计算图中等待反向使用,当层数深、batch 大时,这部分远超参数显存。
核心原理
学霸处理草稿纸的秘诀其实很简单:不是所有中间步骤都值得记。他会把题目分成几段,在每段结束时抄下一个“检查点”结果;其余过程只在草稿上划几下,用完就翻页。AI 的梯度检查点原样复刻了这个思路:前向传播时,不再把所有层的激活都保存在显存里,而是每隔几层才保存一个检查点激活,其余激活算完立即释放。
当反向传播需要某个被丢弃的激活来计算梯度时,怎么办?从最近的检查点开始,重新执行一遍前向传播,把这条路走一遍,顺路就把激活“再生”出来了。这就是“重算”(recompute)——相当于你从上一个关键步骤开始,把中间的演算快速重推了一遍。
激活显存被压缩到一个很小的量级,代价是反向传播时多了几次“回头重推”。但就像学霸宁愿多花两分钟重推一段演算,也不愿向老师再要十张草稿纸——算力换显存,这就是梯度检查点最核心的一笔交易。
- 全量保存普通模式下,前向传播算出每一层的激活后全部留在显存,总共占用 N 层 × 每层 a。
- 检查点保存启用梯度检查点后,只在 K 个检查点位置保存激活,其余层激活用完即释放。
- 压缩比激活显存从 N 份降到 K 份,节省比例约为 N/K 倍。
选取策略
检查点放在哪里,决定了这笔交易划不划算。如果检查点放得太密,每个检查点本身也要占显存,压缩效果大打折扣;如果放得太疏,每段重算路径太长,额外计算又太多。学霸的做法是:把题目按页数均匀分段,每做完固定页数就记一个关键结果。AI 里的默认策略同样是均匀分段。
具体来说,如果有 N 层网络和 K 个检查点,就在第 0、N/K、2N/K……层的位置设置检查点。这样每一段的长度基本相同,重算路径长度均衡,显存和计算的开销被均匀摊开。更高级的做法是“关键层优先”:有的层计算代价小但激活特别大,在它前面放检查点性价比最高;而自适应选择方法则会根据各层实际的内存占用和计算时间,用动态规划找出最优检查点位置。
均匀分段是最稳妥的“存档位”,而关键层优先和自适应选择则是在不同模型结构下寻找更优的存档位。对初学者来说,记住比喻即可:检查点就像书签,每隔差不多厚度的章节夹一个,回头翻书才最省力。
- 检查点集合式子左边的 C 表示所有检查点所在层的编号集合。
- 均匀间隔第 i 个检查点位于 i × N / K 层附近,这样 N 层被 K 个检查点等分成 K 段。
- 下取整因为层号必须是整数,所以用向下取整把浮点位置落到最近的层上。
正确性保障
你可能会担心:重算出来的激活,和原来前向传播算出来的,真的完全一样吗?答案是肯定的。因为每一层的前向计算都是一个确定性函数——同样的输入、同样的权重、同样的算子,无论执行多少次,输出都一模一样。就好比 3 + 5 不管你用笔算还是用计算器,结果永远是 8。
正因为重算结果与原始激活在数值上完全一致,梯度计算也就得到保障。反向传播的梯度是所有层局部导数的乘积链路,链路中的每个“因子”只要数值正确,最后算出的梯度就与普通模式分毫不差。重算只是把链路中缺失的片段重新补齐,并没有改动链路的任何规则。
重推的草稿与原来的草稿内容一致,抄到答题卡上的答案自然一致。因此梯度检查点不会改变训练结果,它改变的只是“草稿纸上的字何时写、何时擦”。
- 损失对第 j 层输入的梯度这是我们最终要恢复的目标,记为 \( \frac{\partial \mathcal{L}}{\partial x_j} \)。
- 来自下一层的梯度它等于第 j+1 层传来的梯度乘以本层函数的导数。
- 重算保证 x_j 正确只要通过检查点重算出的 \( x_j \) 和原始数值相同,上式右端两个因子就都与普通模式一致。
权衡分析
任何交易都要看性价比。梯度检查点用计算换显存,这笔账到底划不划算?我们算一笔量化账。设网络共有 N 层,均匀设置 K 个检查点。普通训练中,前向传播每层算一次(共 N 次层运算),反向传播每层算一次(共 N 次),总计约 2N 次层运算。启用检查点后,反向传播时每个段都要从检查点重新前向算一遍,额外增加约 N-K 次层运算,总计约 3N-K 次。
因此,显存压缩比约为 N/K,而计算开销大约从 2N 增加到 3N,多出约 50%。N/K 越大——即检查点越稀疏——显存省得越多,但重算段越长,计算浪费也越大;N/K 越小,显存节省越有限,但计算开销更接近普通模式。理想情况下,工程上建议 K 取 N 的 1/4 到 1/8,在显存和速度之间取得平衡。
这笔账的结论是:显存省下大头,时间多花一半。当你因为显存不足而无法把 batch 调大、甚至无法启动训练时,这 50% 的时间开销是完全可以接受的——毕竟“能算完”比“算得快”重要得多。
- 普通模式计算量一次前向 N 次加一次反向 N 次,共 2N 次层运算。
- 检查点模式计算量前向仍为 N 次,反向时每个段重算约 N-K 次,再加上反向本身 N 次。
- 增量比例当 K 远小于 N 时,计算量约为普通模式的 1.5 倍。
工程实践
把学霸的做题经验固化成模板,谁都能直接套用。在 PyTorch 中,梯度检查点已经被封装成官方工具:torch.utils.checkpoint.checkpoint 和 checkpoint_sequential。前者包装单个模块或函数,后者直接作用于一个 nn.Sequential 模型,按指定的段数自动切分并插入检查点。
实践中你只需要在模型 forward 里把想要检查点化的模块交给 checkpoint(block, x),剩下的“丢弃激活 + 重算”逻辑全部由框架完成。训练循环几乎不用改动——损失计算、反向传播、优化器更新照旧。唯一要注意的是:启用检查点后,前向传播会用 torch.no_grad() 执行一次,反向时再重新计算,因此要确保模型中没有依赖前向全局状态的随机操作,否则重算结果可能与原来不一致。
把经验做成模板之后,接入成本几乎为零:checkpoint_sequential(model, segments=4, input=x) 一行代码,显存立刻省下一大截。这也是为什么梯度检查点能成为大模型训练框架中的默认选项之一。
- 总层数\( N_{\text{layers}} \) 是网络全部层数,比如 32 层 Transformer 块。
- 分段数\( S_{\text{segments}} \) 是你要把网络分成几段,段数越多检查点越多。
- 每段大小用总层数除以段数并向上取整,得到每段包含的层数。
对比演进
学霸的草稿纸危机,从来不是一张纸能解决的。除了“关键步骤重推”的梯度检查点,工具箱里还有其他省纸技巧:混合精度把钢笔换成细针管笔,让每个字占的墨水更少(数值精度从 32 位降到 16 位,显存直接减半);梯度累积把大题的解题过程拆成几个小问,分几次草稿完成(用小 batch 模拟大 batch);流水线并行则把题目拆给多个同学分工演算(模型分片到多张卡)。这些方法与梯度检查点并不冲突,而是互补的。
实际的大规模训练里,往往同时开启混合精度和梯度检查点,再叠加流水线并行,把显存压力从“单位面积字密度”“关键步骤存档”“多人协作”三个维度同时缓解。未来的趋势也是如此:编译器层面的自动检查点策略、基于运行时显存监控的动态重算决策,正在让“在哪里存档、存几个档”这件事从人工调参走向自动化。
草稿纸不够用的解法从来不止一种。梯度检查点教会我们的核心智慧是:不要急着换更大的纸,先想想哪些字可以晚点再写——这种“延迟计算”的思想,正从显存优化扩展到更多 AI 系统领域。
- 单精度显存32 位浮点数每个参数和激活占 4 字节,记为 \( M_{\text{float}} \)。
- 半精度显存16 位浮点数每个元素占 2 字节,显存需求约为原来一半。
- 组合使用混合精度省掉一半“墨水”,梯度检查点省掉大量“草稿字”,两者叠加效果显著。
AI 应用场景
大语言模型全参数微调
阶段:模型微调(训练)
解决的问题:微调 LLaMA-7B 这类大模型时,仅激活值就可能占满 24GB 显存,batch 只能设为 1,训练又慢又不稳定。传统做法是换更大显存的卡(成本高),而梯度检查点把激活显存压缩到原来的 1/4——1/8,让 batch 可以提到 8 或 16。
具体案例:某算法团队在单张 A100(80GB)上微调 7B 模型,开启 torch.utils.checkpoint 并将检查点间隔设为 4 层后,batch 从 2 提升到 16,训练吞吐提升约 2.5 倍。
长序列 Transformer 预训练
阶段:模型预训练
解决的问题:Transformer 的激活显存随序列长度线性增长,处理 16K 长文本时激活值膨胀到数十 GB。梯度检查点通过分块重算,让长序列训练在有限显存下得以推进。
具体案例:训练一个 1.3B 参数、处理 8192 token 上下文的模型时,开启检查点后单卡即可完成前向和反向,避免使用昂贵的多卡张量并行来“硬扛”显存。
AIGC 扩散模型训练
阶段:生成模型训练
解决的问题:扩散模型 U-Net 在不同分辨率下生成大量中间特征图,显存峰值极高。利用梯度检查点释放中间激活,允许更高的图像分辨率或更大的 batch。
具体案例:某 AIGC 团队在 512×512 分辨率下微调 Stable Diffusion,原本 batch 只能设为 1,启用检查点后 batch 提升到 4,训练速度明显加快。
Python 演示
import torchimport torch.nn as nnimport torch.nn.functional as Ffrom torch.utils.checkpoint import checkpoint torch.manual_seed(0) class MLPBlock(nn.Module): def __init__(self, dim_in, dim_out): super().__init__() self.fc = nn.Linear(dim_in, dim_out) def forward(self, x): return F.relu(self.fc(x)) class Net(nn.Module): def __init__(self, depth=8, use_ckpt=True): super().__init__() self.blocks = nn.ModuleList([MLPBlock(16, 16) for _ in range(depth)]) self.use_ckpt = use_ckpt def forward(self, x): for i, block in enumerate(self.blocks): if self.use_ckpt and i % 2 == 0: x = checkpoint(block, x) else: x = block(x) return x def get_grads(net, x, y): loss = F.mse_loss(net(x), y) net.zero_grad() loss.backward() return [p.grad.clone() for p in net.parameters()] x = torch.randn(4, 16)y = torch.randn(4, 16) plain_grads = get_grads(Net(use_ckpt=False), x, y)ckpt_grads = get_grads(Net(use_ckpt=True), x, y) ok = all(torch.allclose(a, b, atol=1e-6) for a, b in zip(plain_grads, ckpt_grads))print("梯度一致:", ok)print("检查点模式省显存: 偶数层激活不驻留, 反向时重算")
from torch.utils.checkpoint import checkpoint导入 PyTorch 官方梯度检查点工具,它是本次演示的核心 API。x = checkpoint(block, x)用检查点模式运行该模块:前向时释放中间激活,反向时自动重算。self.use_ckpt = use_ckpt用一个开关控制是否启用检查点,方便同一套代码跑出两种模式做对比。plain_grads = get_grads(Net(use_ckpt=False), x, y)普通模式训练一轮,收集各参数梯度作为基准。ckpt_grads = get_grads(Net(use_ckpt=True), x, y)检查点模式训练一轮,收集对应的梯度用于比对。torch.allclose(a, b, atol=1e-6)逐参数判断两组梯度误差是否在 1e-6 以内,证明检查点不改变数值结果。
常见误区
误区好多教程说它会“丢弃激活”,那训练出来的模型精度肯定比普通模式差吧?我是不是不该在生产环境里用它?
正解丢弃的不是信息,而是激活的“驻留权”。重算是严格确定性操作:同一输入、同一权重、同一算子,重算出来的激活与原来逐位一致。梯度由链式法则保证也完全一致。因此模型精度与普通模式等价,只是训练时间略长。
为何常见:“丢弃”这个词具有误导性,让人以为信息被破坏了,实际上只是“暂时没存、需要时再算”。
误区既然检查点是存档,那我把每一层都设成检查点,不就每个中间结果都能快速拿到了吗?显存应该更省吧?
正解每个检查点本身也要保存一份激活,K 越大,需要常驻显存的激活反而越多。而且检查点密了段长变短,重算开销虽然变小,但显存节省也变少。最佳 K 大约在 N/4 到 N/8,而不是越大越好。
为何常见:日常经验中“多备份更安全”的直觉,被错误迁移到了“多存档更省内存”上。
误区我的模型才几百万参数,用不上这种高级技巧,等以后做大模型再说。
正解只要训练时显存受限——无论是想把 batch 调大、把序列调长,还是想在更小的消费级显卡上跑通实验——梯度检查点都能用上。很多中小模型为了提升 batch 也在广泛使用。
为何常见:技术文章总拿几十亿参数模型举例,让中小模型使用者误以为事不关己。
误区开启检查点后显存降了那么多,训练时间好像也没啥变化,这简直是免费的显存外挂。
正解反向传播中每个段都要重新执行一遍前向,计算量约增加 50%。如果只看到显存下降而忽略训练变慢,可能在 deadline 前才发现时间不够。显存与时间之间存在明确的取舍关系。
为何常见:显存数字的变化很直观,而 FLOPs 的增加不直观,容易被忽略。
误区反正激活都能重算出来,我随手在网络的第 2 层、第 7 层设检查点,效果应该差不多。
正解检查点位置决定每个重算段的长度。段长越不均匀,训练越慢;放得不好甚至可能让某些段的重算路径特别长。默认应该均匀分段,并结合各层激活大小与计算代价做调整。
为何常见:不了解重算的计算量由“段长”决定,以为检查点只是一些无关紧要的标记。
面试题
Q1. 梯度检查点主要节省的是哪一部分显存?
- A.模型参数值
- B.训练数据集
- C.中间激活值← 正确答案
- D.优化器状态
答案:C。梯度检查点针对的是前向传播中逐层产生的中间激活。参数、梯度、优化器状态的大小并没有改变,因此节省的核心是激活显存,其他部分基本不受影响。
Q2. 反向传播需要某个已被丢弃的激活值,如何找回它?
- A.随机初始化生成
- B.直接复制检查点
- C.从内存文件读取
- D.重新前向计算← 正确答案
答案:D。被丢弃的激活没有副本,只能从最近检查点开始重新执行前向传播来重建。这是“重算”的本质,也是计算换显存的核心来源。
Q3. 使用梯度检查点训练,主要的额外代价是什么?
- A.更多训练轮次
- B.更多计算时间← 正确答案
- C.更多显存空间
- D.更高数值精度
答案:B。检查点模式在反向传播时要额外执行若干次前向重算,因此计算量上升、训练时间变长。它省的是显存而不是时间,更不涉及精度提升。
Q4. 大模型训练中显存占用最大的通常是哪一部分?
- A.模型参数
- B.优化器状态
- C.中间激活值← 正确答案
- D.损失函数值
答案:C。参数和优化器状态与参数量线性相关,而激活值与层数、batch、序列长度同时相关,在深层大模型中往往膨胀为最大的显存开销,因此成为梯度检查点优化的首要目标。
Q5. 关于梯度检查点,下列说法正确的是哪些?
- A.前向只保存部分激活← 正确答案
- B.重算保证数值等价← 正确答案
- C.检查点越多重算开销越大← 正确答案
- D.计算与显存无需权衡
答案:ABC。前向保存部分激活是检查点的定义;重算是确定性操作,梯度数值等价;检查点越多重算段越短、重算开销越小而非越大,所以 C 错误。计算与显存之间存在明确取舍,D 也错误。
Q6. PyTorch 中梯度检查点工具位于哪个模块?
- A.torch.utils.checkpoint← 正确答案
- B.torch.nn.checkpoint
- C.torch.optim.checkpoint
- D.torch.utils.memory
答案:A。官方实现位于 torch.utils.checkpoint,提供 checkpoint 和 checkpoint_sequential 等接口。其他选项均为不存在或错误位置的模块名。
Q7. 设网络共 N 层,均匀设置 K 个检查点,重算激活带来的额外前向计算层数约为多少?
- A.N 减 K← 正确答案
- B.N 除以 K
- C.K 乘 N
- D.N 的平方
答案:A。每个非检查点层在反向时都可能被重算一次,除检查点层外,总计约 N-K 层需要额外前向计算。因此额外计算量随 K 增大而减小,这也是 K 不能过小的原因。
Q8. 普通前向传播时,网络会默认保存什么供反向使用?
- A.最后一层的输出
- B.模型的参数副本
- C.所有层的中间输出← 正确答案
- D.学习率调节状态
答案:C。普通模式为计算梯度,前向时会把每一层的中间输出都保存到计算图中。这带来完整梯度链路,但也正是这些“默认保存”耗尽了显存。
Q9. 梯度检查点可以与哪些技术互补使用?
- A.混合精度← 正确答案
- B.梯度累积← 正确答案
- C.流水线并行← 正确答案
- D.数据增强
答案:ABC。混合精度压缩每个张量的位宽,梯度累积用小 batch 模拟大 batch,流水并行把模型分到多卡,三者与检查点的优化维度互补,可叠加使用。数据增强与显存优化无关。
Q10. 梯度检查点的核心设计哲学是什么?
- A.用算力换显存← 正确答案
- B.用显存换精度
- C.用数据换时间
- D.用精度换速度
答案:A。检查点模式通过反向重算的额外计算,换取前向激活显存的下降,本质是“算力—显存”之间的杠杆,而不是在精度或数据上做妥协。