激活重计算Activation checkpointing
保留部分输入或检查点，在需要中间结果时重放计算，以计算换内存。
- 核心约束
- 重放须恢复相容的输入与执行状态，生成原计算所需的中间结果。
- 适用问题
- 训练显存受限且部分前向计算可安全重放的计算图。
- 边界与取舍
- 副作用、随机状态和设备迁移会影响重放正确性；与持久化模型检查点不同。
常见操作#
以计算图中的检查点区间为单位；不能仅用层数概括所有模型。
| 操作 | 时间成本 | 条件与说明 |
|---|---|---|
| 前向保存检查点 | 保留边界输入和必要状态 | 丢弃可重算的中间激活，具体保存内容由算子和策略决定。 |
| 反向重算 | 增加相应区间的前向计算 | 非重入实现可在所需激活重建后提前结束；不能统一说成本恰好翻倍。 |
| 随机状态处理 | 保存 / 恢复 RNG 状态的额外工作 | 重算须与原前向语义相容；副作用、设备移动与全局可变状态需检查。 |
时间与空间#
- 时间成本
- 正常前反向计算 + 被重算子图的计算
- 空间成本
- 检查点激活 + 重算期间工作集
节省的是部分激活内存，不自动减少模型参数、梯度和优化器状态。内存与计算取舍取决于分段和执行策略。
工程中的应用#
1 个案例 · 附源码PyTorch 的非重入 activation checkpoint 在反向传播需要中间张量时重新执行对应前向片段，以额外计算换取更低的激活存储需求。
附源码片段 · 内存管理
基础原理参考资料
PyTorch 2.5 · Checkpoint重入 / 非重入差异、重算提前结束和 RNG 保持。
PyTorch 2.5.1 · checkpoint.pycheckpoint 的状态捕获与重算实现。