文章

LLM Activation Checkpointing 生产实战：用 Selective Recompute、非重入模式与 RNG 一致性压低训练显存

大模型训练中，激活值常比参数更早耗尽显存。本文深入探讨选择性重计算、非重入模式、随机数状态一致性及工程门禁策略，帮助团队在保证梯度正确的前提下降低峰值显存，控制训练吞吐损失。

背景：训练显存不只被参数占用

部署推理服务时，团队通常先关注权重、KV Cache 和显存分配器；进入全量微调或持续预训练后，真正先把显存推到上限的经常是激活值。

在前向传播期间，Autograd 需要保存一批中间张量，供反向传播计算梯度。这部分内存通常会随微批大小、序列长度、隐藏维度和网络层数增长。可以用下面的粗略关系理解趋势：

activation_memory ∝ micro_batch × sequence_length × hidden_size × layers

它不是精确公式——注意力、MLP、中间投影、并行策略、数据类型和 Kernel 实现都会改变常数项——但足以解释一个常见现象：模型参数已经通过 FSDP 或 ZeRO 分片，序列长度继续增加时，训练仍然可能因激活值而 OOM。

Activation Checkpointing 的基本思路是：前向时少保存一部分中间激活，反向时重新执行对应前向片段，把显存压力转换成额外计算。它解决的不是”训练状态如何落盘”，也不是故障后的断点恢复；这里的 checkpoint 指的是计算图中的激活检查点。

核心原理：保存边界，反向时重放内部计算

普通训练会保存反向所需的中间张量。启用 Activation Checkpointing 后，一个被包裹的区域通常只保留必要输入；当 backward 需要内部激活时，再次执行该区域的前向逻辑。

这带来三个直接结果：

  1. 峰值显存下降：大量中间激活不必跨越整个前向阶段存活。
  2. 计算量上升：被丢弃的激活必须在反向阶段重新生成。
  3. 正确性约束变严：重放必须与第一次前向在控制流、随机数和设备语义上等价。

因此，工程目标不应写成”把所有 Transformer 层都 checkpoint”。更合理的目标是：

在满足目标微批、序列长度和峰值显存水位的前提下，使重计算开销最小，并证明梯度语义没有变化。

不要把整块重计算当成唯一方案

区域级 Activation Checkpointing

最常见的做法是按 Transformer Block 划分区域，例如每 2 层、4 层或 8 层形成一个 checkpoint region。

区域过大时，反向需要重算大量昂贵算子；区域过小时，边界张量、函数调用和调度开销会增多。实际最优点与模型结构、序列长度、FlashAttention 后端及微批配置有关，不能只凭层数决定。

Selective Activation Checkpointing

PyTorch 的 Selective Activation Checkpointing（SAC） 允许在同一个区域内进一步指定：

  • 哪些算子输出必须保存；
  • 哪些算子优先重计算；
  • 哪些算子可考虑保存后转移到 CPU。

一个实用起点是：优先保存计算昂贵的矩阵乘法和注意力结果，重算便宜的逐元素算子、激活函数和部分归一化辅助计算。这不是固定规则，最终仍需通过 Profiler 和显存快照验证。

PyTorch 官方博客把 SAC 描述为普通 Activation Checkpointing 之上的细粒度策略层：普通 AC 会重算区域内所有操作，而 SAC 可以避免重复执行部分高成本操作。

为什么生产代码应优先非重入模式

PyTorch 当前同时提供重入和非重入两种 checkpoint 实现，并明确推荐显式使用 use_reentrant=False。

非重入模式的工程优势包括：

  • 所需激活已经重建后，可以提前停止后续无用重计算；
  • 支持更完整的 Autograd 使用方式；
  • 能处理关键字参数和更复杂的嵌套输入；
  • 可启用 determinism_check 与 debug，对比原前向和重放张量的形状、类型、设备及算子轨迹。

重入模式会完整重跑整个函数，并且对输入输出 requires_grad、嵌套结构、detach() 和 backward API 存在更多限制。旧项目升级时，不能只把默认参数改掉就结束，还要验证自定义 Autograd Function、梯度 Hook 和模块副作用。

RNG 一致性：Dropout 正确不等于随机数一定正确

Activation Checkpointing 会多执行一次前向片段。如果没有处理随机数状态，Dropout、随机深度、采样路由或其他随机算子可能在重放时生成不同结果，从而改变梯度。

PyTorch 默认会保存并恢复 RNG 状态，使 checkpoint 版本尽量与未 checkpoint 的版本保持确定性。这个过程有额外开销，但在建立正确性基线时不应贸然关闭。

需要特别注意三个边界：

  1. 默认逻辑通常处理 CPU 和一种推断出的设备类型；一个 checkpoint 区域跨越多种设备类型时要单独验证。
  2. 如果函数内部把张量迁移到此前未出现的新设备，框架无法提前保存该设备的 RNG 状态。
  3. 全局随机生成器、自定义 CUDA Kernel 或第三方算子的 RNG 语义可能不受通用机制完整覆盖。

因此，preserve_rng_state=False 只适合在确认区域没有随机操作，或业务明确接受随机轨迹变化后使用。它是性能选项，不应成为默认”优化开关”。

工程落地：从整块重算逐步收敛到选择性策略

第一步：建立未启用重计算的基线

至少记录以下指标，并固定模型、数据批次、Seed、并行度和混合精度配置：

  • max_memory_allocated 与 max_memory_reserved；
  • 每 Step 时延及 Tokens/s；
  • 前向、反向、优化器阶段耗时；
  • Loss、梯度范数和非有限值计数；
  • 长序列与短序列分桶下的峰值显存；
  • 每个 Rank 的显存水位与拖尾时间。

只记录”能否跑起来”不够。Activation Checkpointing 经常让 OOM 消失，却可能把每步时间拉长到无法接受。

第二步：先按完整 Transformer Block 启用

先采用可解释的 Block 粒度，例如只对中后段若干层启用，而不是一次覆盖全部模型。逐步增加覆盖层数，得到显存—吞吐曲线。

NVIDIA Megatron Bridge 的性能指南区分了只重计算注意力相关激活与完整 Transformer Block 重计算，并允许配置重计算层数。官方指南也指出，完整 Block 重计算可显著降低激活显存，但会引入明显额外计算，因此它应被视为强力但昂贵的档位。

第三步：用 Selective Recompute 保留昂贵算子

下面给出一个简化的 PyTorch 示例。具体算子名称会随 PyTorch 和 Attention Backend 版本变化，生产代码应建立版本兼容测试。

from functools import partial
import torch
from torch.utils.checkpoint import (
    CheckpointPolicy,
    checkpoint,
    create_selective_checkpoint_contexts,
)

# 示例：避免重算高成本矩阵乘法。
OPS_TO_SAVE = {
    torch.ops.aten.mm.default,
    torch.ops.aten.bmm.default,
    torch.ops.aten.addmm.default,
}

def policy_fn(ctx, op, *args, **kwargs):
    if op in OPS_TO_SAVE:
        return CheckpointPolicy.MUST_SAVE
    return CheckpointPolicy.PREFER_RECOMPUTE

sac_context = partial(create_selective_checkpoint_contexts, policy_fn)

def run_block(block, hidden_states, attention_mask):
    # 用闭包固定关键字参数，减少模型接口差异。
    def forward_fn(x):
        return block(
            hidden_states=x,
            attention_mask=attention_mask,
        )
    return checkpoint(
        forward_fn,
        hidden_states,
        use_reentrant=False,
        preserve_rng_state=True,
        determinism_check="default",
        context_fn=sac_context,
    )

策略清单不能只看算子名称。需要在目标模型上确认：

  • Attention 是否被融合成不同的 ATen 或自定义算子；
  • torch.compile 是否改变图分区与保存策略；
  • FlashAttention、Transformer Engine 或 Triton Kernel 是否能被策略函数识别；
  • 保存某个输出是否反而增加更大的生命周期内存。

第四步：把”正确性回放”设为发布门禁

建议为同一批固定数据运行三条轨道：

  1. 不启用 Activation Checkpointing；
  2. 启用普通 Block Checkpointing；
  3. 启用 Selective Activation Checkpointing。

对比以下内容：

  • 每步 Loss 差异；
  • 关键层梯度的最大绝对误差、相对误差和余弦相似度；
  • 参数更新后的权重差异；
  • Dropout 开启时多次固定 Seed 回放；
  • 单卡、DDP、FSDP 或张量并行下的结果一致性；
  • 梯度累积多个 micro-step 后的误差是否扩大。

数值容差应根据 BF16、FP16、FP8、算子后端和并行归约顺序制定，不应要求所有配置逐 Bit 相同，也不能只比较最终 Loss。

分布式训练中的额外风险

各 Rank 必须使用相同策略

如果某些 Rank 因条件分支进入 checkpoint 区域，而另一些没有进入，反向阶段的集合通信顺序可能不一致，最终表现为 NCCL 卡死，而不是明确的 Python 异常。

避免在重放区域中执行副作用

以下逻辑应尽量放到 checkpoint 区域之外：

  • 更新全局计数器；
  • 写日志、缓存或队列；
  • 修改模块外部状态；
  • 仅应执行一次的 AllReduce 或 RPC；
  • 依赖当前时间、请求 ID 或非确定性外部输入的分支。

PyTorch 官方文档明确警告：如果 backward 重放时的函数调用与原前向不同，可能报错，也可能产生静默的错误梯度。

Pipeline Parallel 要按 Stage 评估

不同 Pipeline Stage 的层数、激活体积和气泡比例可能不同。统一设置”每个 Stage 重算 N 层”并不一定合理。应分别测量各 Stage 峰值显存和拖尾时延，再调整重计算覆盖范围。

如何选择策略

训练状态优先方案原因
只差少量显存即可提升微批少量 Block Checkpointing配置简单，易于验证
长序列导致注意力激活占主导注意力选择性重计算通常能针对主要内存项
完整 Block 重算吞吐损失过大SAC 保留 Matmul/Attention避免重复昂贵算子
FSDP 后参数内存已低、激活仍高AC 与激活 Offload 对比压测瓶颈已从模型状态转向激活
Dropout 或随机路由较多保留 RNG 状态并做梯度对照防止重放随机轨迹变化
模型含复杂嵌套输入和自定义梯度非重入模式兼容性和诊断能力更完整

常见误区

误区一：显存下降比例可以直接照搬论文

不同模型、序列长度、并行策略和 Kernel 后端差异很大。关于选择性重计算的经典研究在特定大规模 GPT 训练配置中报告了显著的激活内存和执行效率收益，但这些结果不能直接视为任意微调任务的预期值。

误区二：只看 GPU Utilization 判断收益

重计算可能让 GPU 更忙，但 Tokens/s 下降。需要同时观察每步时延、有效 Token 吞吐、MFU、峰值显存和训练总成本。

误区三：关闭 RNG 保存就是免费提速

如果区域内存在 Dropout 或其他随机操作，关闭 RNG 状态保存会改变重放路径。必须先证明梯度差异处于可接受范围。

误区四：Activation Checkpointing 与持久化 Checkpoint 是同一件事

前者是计算—显存权衡，发生在每个训练 Step 的前向和反向之间；后者是故障恢复和训练状态持久化，发生在存储与恢复链路。两者应分别设计和监控。

误区五：重放区域越大越省事

区域越大，往往越省显存，但计算和调试成本也越高。生产配置应寻找满足显存目标的最小重计算集合。

上线检查

  • 已固定模型版本、数据批次、Seed、并行度和混合精度建立基线。
  • 显式设置 use_reentrant=False，并验证自定义 Autograd 和梯度 Hook。
  • 已确认重放区域内不存在不可重复副作用。
  • Dropout、随机路由和自定义 RNG 已做确定性回放。
  • 已按短、中、长序列分别测量峰值显存和 Tokens/s。
  • 已对比无 AC、Block AC、Selective AC 三条轨道的 Loss 和关键层梯度。
  • 各分布式 Rank 使用一致的 checkpoint 策略和控制流。
  • 已验证 torch.compile、FlashAttention 和自定义 Kernel 下的策略兼容性。
  • 已设置 OOM、非有限梯度、Step 时延和吞吐退化告警。
  • 保留一键关闭 Selective 策略并回退到已验证 Block 配置的开关。

参考资料

  1. PyTorch torch.utils.checkpoint 文档：https://docs.pytorch.org/docs/2.13/checkpoint.html
  2. PyTorch Foundation，Current and New Activation Checkpointing Techniques in PyTorch：https://pytorch.org/blog/activation-checkpointing-techniques/
  3. NVIDIA Megatron Bridge Performance Tuning Guide：https://docs.nvidia.com/nemo/megatron-bridge/latest/performance-guide.html
  4. Korthikanti et al.，Reducing Activation Recomputation in Large Transformer Models：https://arxiv.org/abs/2205.05198

常见问题

Activation Checkpointing 会改变模型收敛结果吗？
按设计不应改变，但前向重放必须与原前向等价。随机数、全局变量、设备迁移、副作用和非确定性 Kernel 都可能造成差异，因此必须进行 Loss、梯度和参数更新对照。
为什么优先使用 use_reentrant=false？
PyTorch 推荐非重入实现。它支持提前停止无用重计算、关键字参数、嵌套结构和更完整的 Autograd 场景，同时提供 determinism_check 与 debug 等诊断能力。
Selective Recompute 是否总比完整 Block 重算快？
不一定。它提供更大的策略空间而非自动保证。只有在目标模型、序列长度、并行方式和算子后端上实测显存—吞吐 Pareto 前沿，才能确定最优保存与重算集合。
特效配置