首页优化地图方法索引基础概念
返回技术汇总表 ↗

M037 · 独立详解

激活重计算 / Checkpointing

保存少量边界,反向需要中间结果时重新计算,用额外计算换取激活容量。

内存与数据搬运调度与并行详解已编写
视频讲解待制作

这篇文章的小视频尚未制作;可以先阅读下方文字详解。

本页目录

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 等随机操作需在重算时保持相容的随机状态。区域内改变全局状态、执行不可重复副作用或依赖不同控制流,可能让重算与原前向不一致。

依据:PyTorch · 激活 Checkpoint(main) ↗ · PyTorch · 选择性激活重计算 ↗

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 · 最小接入与一致性检查

PyTorch 接入片段 · 假设 block 与 hidden 已定义;并行组合需另核对
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 涉及随机数一致性,也有管理成本;设备类型变化等特殊情况应查文档。

依据:PyTorch · 激活 Checkpoint(main) ↗

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 文档时,实际接口以安装版本为准。

代码为接入片段或教学伪代码;本页没有声称完成 NVIDIA / 昇腾硬件实测。

← 返回总表中的 M037继续阅读:训练状态 / 激活卸载与回迁 →
AI Infra 知识库 · 总览 → 独立讲解网站更新 2026-10-08
搜索知识库

搜索方法与基础概念 · Esc 关闭