梯度检查点是通过只保存部分中间激活值、反向时重算前向来节省显存的技术，能降低40%~60%显存但增加15%~30%训练时间，要求模块前向可重入且无副作用。
梯度检查点是什么，为什么能省显存
梯度检查点（torch.utils.checkpoint.checkpoint）不是“不存梯度”，而是“不存中间激活值”。反向传播时需要前向计算的中间结果来算梯度，常规训练会把所有层的输出全存着，显存占用和网络深度线性增长；而检查点只存部分层的输入，反向时临时重跑对应前向——用时间换空间。
典型节省比例：ResNet-50 训练 batch size 可从 16 提到 32，ViT-L 类模型显存常降 40%~60%。但注意：它只对**前向可重入、无副作用**的模块有效，比如不能包裹含 nn.Dropout（训练态随机行为不可复现）或修改全局状态的自定义层。
怎么加 checkpoint，最简安全写法
别直接套整个模型，先从单个 nn.Sequential 或自定义 forward 块开始。PyTorch 官方推荐方式是用 checkpoint.checkpoint 包裹函数调用，而不是用装饰器（后者容易隐式捕获非 tensor 参数）。
- 必须确保被包裹函数只接收
Tensor参数，且不依赖闭包变量（如self.training） - 若模块含
training切换逻辑（如Dropout），改用checkpoint.checkpoint_sequential或手动拆分 +torch.no_grad()重算 - 示例：对 Transformer 层列表做检查点
from torch.utils.checkpoint import checkpoint
def custom_forward(x, layer):
return layer(x)
# 替换原循环：x = layer(x)
x = checkpoint(custom_forward, x, layer)
常见报错和绕过方法
RuntimeError: Trying to backward through the graph a second time：说明 checkpoint 内部用了被复用的 Tensor（比如共享 embedding），或者你在检查点外又对同一张量调了 backward()。根本原因是计算图被意外保留。
让 macOS 自动化像说话一样简单 —— 用代码控制你的 Mac，无需学习 AppleScript 系统要求 macOS 14+ (Sonoma) Python 3.9+ 需要开启：辅助功能 + 屏幕录制权限
- 确认被 checkpoint 的函数返回值没被其他分支引用（尤其注意 loss 计算前的中间变量）
- 避免在 checkpoint 块里调用
.detach()或.item()—— 这些操作会切断梯度流，导致反向失败 -
torch.cuda.amp混合精度下，务必让checkpoint在autocast上下文内，否则可能因 dtype 不匹配报Inconsistent type
性能与精度影响要实测，不能默认“安全”
加了 checkpoint 后训练变慢是常态，通常慢 15%~30%，因为重算前向。更隐蔽的问题是数值误差：FP16 下多次重算可能放大舍入误差，个别任务（如长序列语言建模）会出现 loss 震荡或收敛变差。
- 建议先关掉
torch.backends.cudnn.benchmark=True，避免不同次前向触发不同卷积算法，加剧不一致 - 验证时务必禁用 checkpoint（设
use_checkpoint=False），否则 eval 结果不准 - 如果模型含
BatchNorm，注意其 running stats 在 checkpoint 块内不会更新——得把 BN 层移出检查点范围，或改用SyncBatchNorm+ DDP 显式同步
真正麻烦的从来不是加不加那行 checkpoint，而是哪些子模块能进、哪些必须剔出来、以及重算带来的随机性是否破坏了你的训练稳定性。
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
