PyTorch torch.utils.checkpoint 怎样用重计算换显存

训练模型时，显存紧张不一定要立刻减小 batch size。torch.utils.checkpoint 提供了一种明确的交换：前向阶段少保存一部分中间激活，反向传播需要它们时再重新计算，因此通常用更多计算时间换取更低的激活显存。它不是把参数压缩，也不会自动保证吞吐更高。

官方地址：https://docs.pytorch.org/docs/stable/checkpoint

最稳妥的用法是只包住计算边界清晰、可以重复执行的前向片段，并显式写出 use_reentrant=False。启用后要同时观察显存、吞吐、梯度一致性和随机行为，不能只看显存下降。

先判断是否适合用重计算换显存

普通训练会保留前向产生的中间激活，等反向计算梯度时使用。activation checkpointing 会让指定函数内部的部分激活不再长期保留，反向阶段重新调用这段函数。适合它的场景通常是模型较深、激活占用明显、计算资源还有余量；如果瓶颈本来就是计算吞吐，重算可能让训练更慢。

先把边界说清楚：checkpoint 这里指激活检查点，不是把模型参数保存到文件的 checkpoint。建议先记录未改造版本的峰值显存、每步耗时和一个固定小批次的损失，再决定是否引入。

用非重入实现包住计算段

把待重算的前向逻辑收敛到一个函数中，然后把输入张量传给 checkpoint。新代码显式指定 use_reentrant=False，便于使用关键字参数、与更完整的 autograd 场景配合，也避免依赖默认值变化。

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

class CheckpointedBlock(nn.Module):
    def __init__(self, width: int) -> None:
        super().__init__()
        # 让前向片段保持纯粹，便于反向阶段安全重算。
        self.block = nn.Sequential(
            nn.Linear(width, width),
            nn.GELU(),
            nn.Linear(width, width),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 非重入实现会在需要梯度时重算片段，减少中间激活的长期保存。
        return checkpoint(
            self.block,
            x,
            use_reentrant=False,
        )

model = CheckpointedBlock(1024)
sample = torch.randn(8, 1024, requires_grad=True)
loss = model(sample).square().mean()
# 反向阶段可能再次执行 self.block，因此要把它当作可重复的前向函数。
loss.backward()

这个示例只表达调用关系，不代表某个具体模型一定会节省相同数量的显存。真正落地时，优先选择重复出现且没有外部副作用的 block；不要把写文件、修改全局计数器或依赖不可重复随机状态的逻辑直接塞进重算函数。

处理随机状态与输入边界

如果片段里包含 dropout 等随机操作，重算时的随机状态会影响结果。PyTorch 默认会保存并恢复 CPU 及相关设备的随机状态，以尽量让 checkpoint 版本与普通前向保持一致，但保存状态本身也有运行时成本。若业务不要求这种确定性，可以评估 preserve_rng_state=False，但这必须成为有记录的取舍，而不是为了跑得快随手关闭。

还要检查三类边界：第一，函数在前向和重算时看到的输入必须一致；第二，闭包捕获的外部变量不能在两次调用之间悄悄变化；第三，函数不能依赖只执行一次的副作用。官方文档也提醒，如果反向重算时的函数调用与前向不同，结果可能报错，也可能产生不易察觉的错误梯度。

def forward_segment(x: torch.Tensor, training_scale: float) -> torch.Tensor:
    # 把影响数值的配置作为显式参数，避免重算时读取漂移的全局变量。
    return torch.nn.functional.dropout(
        x * training_scale,
        p=0.1,
        training=True,
    )

# preserve_rng_state 默认开启；只有确认随机性取舍后才显式关闭。
out = checkpoint(
    forward_segment,
    sample,
    1.0,
    use_reentrant=False,
    preserve_rng_state=True,
)

参数边界与生产审查

debug=True 适合定位原始前向与重算之间的算子差异，且只支持非重入实现；它不是性能开关。early_stop=True 是非重入实现的重算控制，默认会在已经得到反向所需张量后尽早停止，关闭它可能只适合有明确兼容性原因的场景。context_fn 则用于为前向和重算提供成对的上下文管理器，应当把它当成审查边界来维护。

上线前可以按下面的顺序记录结果：

• 显存：固定 batch、输入和设备，比较峰值显存是否达到目标。
• 计算：记录每步耗时和吞吐，确认重算成本在可接受范围内。
• 数值：用固定随机种子比较损失和梯度趋势，重点关注含 dropout 的片段。
• 边界：确认所有输入都通过函数参数传入，删除依赖可变全局状态的隐式读取。
• 回滚：保留一个不使用 checkpoint 的配置开关，出现异常时能恢复原路径。

这套清单比单独追求“显存越低越好”更可靠：checkpoint 的价值是让训练约束变得可管理，而不是让一次配置覆盖所有模型结构。

常见问题

checkpoint 会让参数显存也一起减少吗？

它主要针对前向激活的保存方式，参数、梯度和优化器状态仍然有各自的显存成本。若峰值来自优化器状态，应该单独评估混合精度、分片或优化器方案。

为什么开启后训练变慢？

因为部分前向计算会在反向阶段再次执行。显存压力下降与计算时间增加是同一项交换的两面，应使用固定工作负载比较端到端吞吐。

可以把任意函数传给 checkpoint 吗？

不能简单按“能调用”判断。函数应尽量无副作用、输入明确且前向与重算等价；涉及可变全局状态、外部 I/O 或不稳定随机行为时，要先重新设计边界。