PyTorch的torch.utils.checkpoint.checkpoint通过前向时丢弃中间激活、反向时重新计算来节省显存，将显存复杂度从O(L)降至O(√L)，但以15%–30%速度下降为代价；需确保纯函数式调用、合理分段、正确设置use_reentrant=False并避免嵌套或大操作。
PyTorch的torch.utils.checkpoint.checkpoint为什么能省显存？
它通过用时间换空间：在前向传播时只保存部分中间激活值，反向传播时重新运行对应子图来重算梯度，跳过对大部分中间张量的持久化存储。典型场景下可将显存占用从 O(L) 降到 O(√L)（L 是层数），但训练速度会下降 15%–30%。
关键限制是：被 checkpoint 包裹的子模块必须是纯函数式调用，不能有 in-place 操作、不可导控制流或依赖外部状态（比如 self.training 的条件分支若影响计算图结构，就会出错）。
怎么正确包裹一个nn.Module子类？
不能直接对整个模型调用 checkpoint，必须拆到「可重计算」的最小粒度——通常是单个 Transformer 层或残差块。常见错误是把含 dropout 或 layer norm 的模块整体包进去，结果反向时因 RNG 状态不一致而梯度异常。
- 推荐做法：自定义
forward方法，在内部用torch.utils.checkpoint.checkpoint包裹核心计算逻辑（如self.attn(x)+self.mlp(x)），把 norm 和 residual 这类轻量操作留在外面 - 必须传入
use_reentrant=False（PyTorch ≥ 2.0），否则可能触发RuntimeError: Trying to backward through the graph a second time - 输入张量需设
requires_grad=True，否则 checkpoint 会静默跳过重计算
示例：
立即学习“Python免费学习笔记（深入）”；
def custom_forward(x):
return self.mlp(self.attn(x))
<p>x = checkpoint(custom_forward, x, use_reentrant=False)
x = x + residual # residual 在外侧
使用FSDP或DeepSpeed时还能开checkpoint吗？
可以，但顺序很重要：checkpoint 必须在 FSDP 包装之后应用，否则 FSDP 的参数分片逻辑和 checkpoint 的重计算会冲突，导致 RuntimeError: input is not contiguous 或梯度为 None。
DeepSpeed 的 stage 3 默认已集成 checkpoint，手动再开容易重复触发重计算；如果要用自定义 checkpoint，需关闭 DeepSpeed 内置选项："activation_checkpointing": {"partition_activations": false}，再自行注入。
混合精度训练中，amp 自动处理 checkpoint 内部的类型转换，但要注意：若 checkpoint 函数内有 .half() 显式转类型，会导致重计算时类型不一致，引发 Expected all tensors to be on the same device and have the same dtype。
为什么有时开了checkpoint反而OOM？
最常见原因是 checkpoint 区域过大或嵌套太深：比如把整个 encoder 一次性包进一个 checkpoint，重计算时临时显存峰值反而更高；或者在 checkpoint 内又调用了另一个 checkpoint，形成递归重计算，显存碎片加剧。
调试建议：
- 用
torch.cuda.memory_summary()对比开启前后各阶段的显存分配模式 - 逐层启用 checkpoint，从最后一层开始往前加，观察显存拐点
- 避免 checkpoint 含有 large embedding lookup 或 broadcast 操作（如
x[:, None, :] * mask），这类操作重算代价高且易爆显存
真正省显存的关键不是“开多少”，而是“在哪开、开多细”——模型越深，越要拆成小段；显存瓶颈在 batch size 上时，checkpoint 效果有限，该调 gradient_accumulation_steps 还得调。
