01 · 反向为什么需要前向中间结果
链式求导不仅需要最终输出。矩阵乘法的参数梯度需要输入,非线性与归一化的反向也需要相应的中间信息。默认自动微分会保存一批 Tensor,直到对应反向完成才能释放;深层网络、长序列和大 microbatch 会抬高激活峰值。
激活重计算主要处理这些中间结果,不直接减少参数与 Adam 状态。这里的 checkpoint 是计算图上的检查点,与定期保存模型文件用于故障恢复是两个概念。
依据:PyTorch · 激活 Checkpoint(main) ↗ · Training Deep Nets with Sublinear Memory Cost ↗
02 · 前向少保存,反向重新产生
forward: save_boundary(x) -> block(x) -> output
do_not_keep_selected_intermediates()
backward: reload_boundary(x) -> recompute(block, x)
use_recomputed_values_for_gradients()可以按一个 Transformer Block、若干层一段,或选择具体操作保留 / 重算。粗粒度省下更多激活,也可能重算昂贵 GEMM;细粒度策略更复杂,但可保留昂贵操作的结果而重算便宜操作。不同实现不一定完整重跑整个区域。
Dropout 等随机操作需在重算时保持相容的随机状态。区域内改变全局状态、执行不可重复副作用或依赖不同控制流,可能让重算与原前向不一致。
03 · 十二层网络的教学容量模型
自设极简模型:12 层,每层边界和所需内部激活各抽象为 A=64 MiB;普通保存按 12A≈768 MiB 估计。把每 3 层作为一段,只保存约 4 个段输入;反向一次重算一段,再临时保存该段约 3 份内部激活,估计 (4+3)A≈448 MiB。
这只是说明“永久保存减少、临时重算仍占内存”。真实层内有多个 Tensor,边界输入与输出可能重叠,最后一段、Loss 和已有梯度也影响峰值,因此不能把 448 MiB 当作实际框架结果。
| 项 | 普通保存 | 按段重算(上述假设) |
|---|---|---|
| 长期保存 | 约 12A | 约 4A 段边界 |
| 反向前额外计算 | 无对应重跑 | 重算被丢弃的区域 |
| 峰值来源 | 累积中间结果 | 边界 + 当前段临时结果 + 其他状态 |
再假设未优化前向 30 ms、反向 60 ms,若需要重算一整次等量前向,单步约从 90 增至 120 ms,而不是必然加速。省下的内存若能换取更有效批量、减少卸载,才可能提高最终吞吐。
04 · 最小接入与一致性检查
from torch.utils.checkpoint import checkpoint
for block in model.layers:
hidden = checkpoint(
block, hidden,
use_reentrant=False,
preserve_rng_state=True
)显式选择 use_reentrant,避免依赖不同版本的默认行为。非重入实现可以在已重建所需 Tensor 后提前停止部分重算。preserve_rng_state 涉及随机数一致性,也有管理成本;设备类型变化等特殊情况应查文档。
05 · 如何选择重算位置
- 先用内存 / profiler 确认激活是主要占用;参数状态很大时先考虑分片。
- 从少量 Block 开启,比较输出、梯度和 Loss,再逐步扩大范围。
- 分别记录峰值内存、前向、反向和重算时间;比较同有效批量的吞吐。
- 与 FSDP / PP 组合时,检查重算是否触发额外参数汇集,是否改变激活驻留。
FlashAttention 的反向已经重算部分 Attention 中间量;在它外面再做整块 checkpoint 可能叠加重算。二者可以组合,但应以实际峰值与耗时衡量,不能简单相加两份宣传收益。
06 · NVIDIA 与昇腾怎样对照
重计算是自动微分和训练运行时策略,两平台都应检查框架实现和并行组合。NVIDIA 侧可以从 PyTorch checkpoint 阅读;昇腾侧可检查 MindSpeed 的重计算配置及其版本支持。比较重算区域、RNG 处理、通信次数和相同有效批量下的时间。
依据:PyTorch · 激活 Checkpoint(main) ↗ · MindSpeed Core · 特性清单 ↗
07 · 自查与小实验
- 画出三个 Block 的前反向,把每个 Tensor 的最后使用时刻标出来。
- 在带 Dropout 的小模型上比较重算开关后的梯度,并观察随机数状态的作用。
- 比较只重算 Attention、只重算 MLP 和整 Block 的时间 / 内存取舍。
流程位置与关联方法
这个方法可能出现在多个位置;计算线包含模型计算与算子实现,执行线包含运行时、内存和通信。
资料与核查范围
原理依据以原始论文和官方文档为主;算例与时间线是教学推导。文档核查日期:2026-10-10。引用 main / latest 文档时,实际接口以安装版本为准。
- MindSpeed Core · 特性清单 ↗
- PyTorch · Performance Tuning Guide ↗
- PyTorch · 激活 Checkpoint(main) ↗
- Training Deep Nets with Sublinear Memory Cost ↗
- PyTorch · 选择性激活重计算 ↗
代码为接入片段或教学伪代码;本页没有声称完成 NVIDIA / 昇腾硬件实测。