Gradient Checkpointing 通过舍弃中间激活值、反向传播时重算前向来节省显存，显存降至约1/2～2/3，但训练时间增加约33%；需显式集成，仅对继承nn.Module且定义forward的模块有效。
GradientCheckpointing 为什么能省显存
显存占用主要来自两部分：模型参数 + 梯度 + 中间激活值。其中，前向传播中每一层输出的激活值（activation）在反向传播时必须复用，PyTorch 默认会把它们全保留在显存里。Gradient Checkpoint 的核心思路是：不存中间激活，而是反向时重算一次前向。代价是多一次前向计算（约增加33%训练时间），但显存可降为原来的 ~1/2～2/3。
它不是“自动开启”的功能，必须显式插入到模型结构中；也不能套在任意模块外就生效——只对继承 torch.nn.Module 且定义了 forward 的子模块有效。
怎么给自定义 Module 开启 gradient checkpoint
最常用方式是用 torch.utils.checkpoint.checkpoint 包裹前向逻辑，但注意：它要求被包裹函数必须支持 torch.Tensor 输入，且不能有不可导的控制流（如 if x > 0: 这类依赖 tensor 值的分支）。
典型写法：
from torch.utils.checkpoint import checkpoint
<p>class MyBlock(torch.nn.Module):
def <strong>init</strong>(self):
super().<strong>init</strong>()
self.linear1 = torch.nn.Linear(1024, 1024)
self.linear2 = torch.nn.Linear(1024, 1024)</p><pre class="brush:php;toolbar:false;">def forward(self, x):
# 不直接调用 linear1(x)，而是用 checkpoint 包裹整个子前向
def custom_forward(x):
x = self.linear1(x)
x = torch.relu(x)
x = self.linear2(x)
return x
return checkpoint(custom_forward, x)
关键点：
-
checkpoint第一个参数必须是纯函数（无 self、无属性访问），所以要把子模块提前绑定好，或改用闭包捕获 - 输入
x必须是torch.Tensor，不能是字典、tuple 等复合结构（除非你手动解包再打包） - 如果 forward 里有 dropout 或 batch norm，要注意它们在重算时会重新采样——可能影响训练稳定性（尤其 dropout rate 高时）
用 torch.utils.checkpoint.checkpoint_sequential 批量处理 Sequential
如果你的模型主干是 torch.nn.Sequential（比如 Transformer 的 encoder layer 堆叠），可以用更简洁的方式：
from torch.utils.checkpoint import checkpoint_sequential <h1>layers 是一个 nn.ModuleList 或 list of modules</h1><h1>segments=2 表示把 layers 分成 2 段，每段内部不保存激活，只存段首输入</h1><p>output = checkpoint_sequential(layers, segments=2, input) </p>
这个函数本质是把模块列表切片后，对每一段调用 checkpoint。好处是不用改每个 block 的代码；缺点是切分粒度粗，无法跳过某些轻量层（比如 LayerNorm），而且要求所有模块都满足 checkpoint 兼容性。
常见坑：
-
segments数不能大于模块数，否则报错ValueError: too many sections - 如果某一层用了
torch.no_grad()或torch.inference_mode()，checkpoint 会失效（因为重算时也进不了该上下文） - 和
torch.compile(model, mode="reduce-overhead")同时用可能冲突，建议先关 compile 再测 checkpoint 效果
HF Transformers 里怎么开 Gradient Checkpoint
Hugging Face 的 transformers 模型基本都支持原生开关，比手写 wrapper 更稳：
model = AutoModel.from_pretrained("bert-base-uncased")
model.gradient_checkpointing_enable() # 就这一行
# 或初始化时传参
model = AutoModel.from_pretrained("bert-base-uncased",
gradient_checkpointing=True)
它背后做了几件事：
- 自动把各 layer 模块注册为可 checkpoint 的子模块
- 重写了
forward，在合适位置插入checkpoint调用 - 禁用掉 layer norm/dropout 在重算时的随机性（比如固定 dropout mask）
但要注意：不是所有 HF 模型都完整支持。比如某些老版本的 OPTForCausalLM 可能漏掉 decoder 的 final layer；遇到 RuntimeError: Trying to backward through the graph a second time，大概率是某层 forward 被意外多次调用，得查 model.config 的 gradient_checkpointing 是否真正透传到了子模块。
显存节省效果高度依赖模型结构和序列长度。短文本（
Python免费学习笔记（深入）：立即使用
在学习笔记中，你将探索 Python 的核心概念和高级技巧！
