梯度检查点：用计算换显存的经典套路
说起梯度检查点（Gradient Checkpointing），我脑子里第一个蹦出来的词就是「时间换空间」。说白了，就是拿计算量去换显存。你想想看，大模型训练最头疼的是什么？显存不够啊！特别是当你把batch size往上提的时候，显存就像漏气的气球一样哗哗往下掉。
我个人习惯把梯度检查点叫做「选择性遗忘」。为什么这么说？因为它的核心思想就是：在前向传播时，只保留一部分中间激活值，剩下的扔掉；反向传播时，再重新计算那些被扔掉的激活值。嗯，听起来有点绕，我们一步步拆解。
原理详解：为什么能省显存？
先说说正常情况。你训练一个深度网络，前向传播时每一层都会产生中间激活值（activation）。这些激活值得存着，因为反向传播时要用它们来计算梯度。一个100层的Transformer，中间激活值能吃掉你80%以上的显存。
梯度检查点的做法很粗暴：我不全存了，我只存几个关键节点的激活值。比如每10层存一个检查点。反向传播时，需要某个中间层的激活值？好，我从最近的检查点开始，重新做一次前向计算，把需要的激活值算出来。
显存节省效果：理论上，如果网络有N层，每k层设置一个检查点，显存占用可以从O(N)降到O(N/k + k)。实际项目中，我见过把显存从24GB压到12GB的情况，代价就是训练时间多了20%-30%。
为什么会这样？因为反向传播时重新计算激活值，相当于多做了一次前向传播。你想想看，本来一次前向一次反向，现在变成了一次半前向一次反向。计算量上去了，但显存下来了。
torch.utils.checkpoint 实现
PyTorch官方提供了现成的接口，就在torch.utils.checkpoint里。我最常用的是checkpoint_sequential和checkpoint两个函数。
先看个简单的例子：
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.layer1 = nn.Linear(1024, 1024)
self.layer2 = nn.Linear(1024, 1024)
self.layer3 = nn.Linear(1024, 1024)
def forward(self, x):
# 对layer2使用梯度检查点
x = self.layer1(x)
x = checkpoint(self.layer2, x) # 这里只保存输入x，不保存中间激活
x = self.layer3(x)
return x注意看，checkpoint函数接收一个可调用对象和它的输入。在前向传播时，它只保存输入x，不保存layer2内部的中间激活。反向传播时，它会用保存的x重新计算layer2的前向结果，然后算梯度。
我曾经踩过的坑：使用checkpoint时，输入张量不能有requires_grad=False的情况，否则会报错。另外，如果你的模型里有dropout或batch norm，要小心处理——checkpoint会改变这些层的随机种子行为，可能导致训练不稳定。
对于序列模型，checkpoint_sequential更方便：
from torch.utils.checkpoint import checkpoint_sequential
# 假设model是一个包含多个层的Sequential
segments = 4 # 分成4段，每段末尾设置检查点
output = checkpoint_sequential(model, segments, input_tensor)分段策略与性能调优
这里有个关键问题：检查点该设多少个？设多了，计算量太大；设少了，显存省不下来。
我总结了一个经验公式：
| 分段数 | 显存节省 | 计算开销 | 适用场景 |
|---|---|---|---|
| 2-4段 | 30%-50% | 10%-15% | 显存刚好够用，想留点余量 |
| 4-8段 | 50%-70% | 20%-30% | 显存严重不足，需要大幅压缩 |
| 8段以上 | 70%+ | 40%+ | 极端情况，比如单卡跑大模型 |
我个人习惯的做法是：先跑一次profile，看看每层的显存占用。把显存占用大的层挑出来，在这些层前后设置检查点。比如Transformer里，attention和FFN的中间激活最吃显存，我会在每层Transformer block的入口设置检查点。
还有一个调优技巧：不要对所有层无差别使用checkpoint。有些层计算量小但显存占用大（比如embedding层），用checkpoint很划算；有些层计算量大但显存占用小（比如简单的激活函数），用checkpoint反而得不偿失。
避坑指南：我曾经在一个50层的ResNet上，给所有层都加了checkpoint，结果训练时间翻了2倍，显存只省了30%。后来改成每5层一个检查点，训练时间只多了15%，显存省了60%。所以，分段策略一定要根据实际模型来调，别无脑全上。
与DDP/FSDP的兼容性
这个问题问的人特别多。梯度检查点和分布式训练框架能不能一起用？答案是：能，但要注意细节。
先说DDP（Distributed Data Parallel）。DDP的核心是梯度同步，它不关心你的前向传播是怎么做的。所以checkpoint和DDP天然兼容，你只需要在模型forward里正常使用checkpoint就行。不过要注意一点：DDP的梯度同步是在反向传播之后触发的，而checkpoint的反向传播会多一次前向计算，这会导致DDP的同步时机稍微延后，但不会影响正确性。
再说FSDP（Fully Sharded Data Parallel）。FSDP会把模型参数分片到多个GPU上，前向传播时动态收集参数。这里有个坑：checkpoint的反向传播会重新执行前向计算，这意味着FSDP需要重新收集参数。如果参数收集的开销很大，可能会抵消checkpoint带来的显存收益。
我建议的做法是：
- 和DDP一起用：直接加checkpoint，不需要额外配置。我试过在8卡DDP上跑GPT-2，checkpoint把每卡的显存从32GB降到了18GB，训练速度只慢了12%。
- 和FSDP一起用：建议把checkpoint的粒度调粗一点，减少参数收集次数。比如原来每层一个检查点，改成每两层一个。另外，FSDP的
reshard_after_forward参数要设为True，这样反向传播时能及时释放参数显存。
实际项目经验：我在一个70B参数的大模型训练任务中，同时使用了FSDP和梯度检查点。FSDP负责参数分片，checkpoint负责减少中间激活。最终在256张A100上，每卡显存控制在38GB左右（总显存80GB），训练效率达到了理论峰值的85%。
嗯，这里还要提一个容易被忽略的点：checkpoint和混合精度训练（AMP）的配合。如果你用了AMP，checkpoint保存的输入张量是fp16的，反向传播时重新计算也是fp16的。这没问题，但要注意梯度缩放——checkpoint重新计算的前向传播不会触发梯度缩放，所以不会影响AMP的正常工作。
最后说一句，梯度检查点不是银弹。如果你的模型本身显存占用就不大，或者你的batch size已经很小了，用checkpoint反而会拖慢训练速度。我一般只在显存成为瓶颈时才考虑它——说白了，就是当你发现「显存不够，但计算资源还有富余」的时候，梯度检查点就是你的最佳选择。