激活检查点：用时间换空间的显存压缩技术

什么是激活检查点

在深度学习中，激活检查点 是一种以计算时间换取显存空间的技术。训练深层网络时，前向传播产生的中间激活值通常会全部保留在显存中，供反向传播计算梯度使用。随着模型规模增大，这些激活值会占用惊人的显存，导致超出硬件限制。激活检查点的核心思路是：不保存所有中间激活，只保留部分关键层（检查点），反向传播时再重新计算被丢弃的中间结果。这一机制能让模型在有限的显存下训练更大的网络或使用更大的批次。

为什么需要激活检查点

现代神经网络动辄包含数千万甚至数千亿参数，同时输入的批量大小和序列长度也在增长。前向传播中每一层的输出都需要留在显存中，等待反向传播使用。以 Transformer 为例，激活显存占用公式大致为：

激活记忆体 ≈ 层数 × 批量大小 × 序列长度 × 隐藏维度 × 精度字节数

当模型规模膨胀时，激活值常成为显存瓶颈。激活检查点通过牺牲一部分计算（重算）大幅降低显存峰值，使得：

• 在单 GPU 上训练原本放不下的模型
• 使用更大的 micro-batch 提高训练稳定性
• 在显存受限的边缘设备上进行微调

核心原理：时间换空间

反向传播计算某一层的梯度，需要该层的输入激活值和下一层传来的梯度。标准做法是前向时保存全部激活，反向直接取用，时间与显存成正比。激活检查点的策略则不同：

1. 划分检查点段：将网络按层或模块切分成若干个段，每段只在末尾保存一份激活作为检查点。
2. 前向传播：段内各层正常计算，但除检查点外不保存中间激活，只输出最终激活。
3. 反向传播：反向到达某个段时，利用该段的输入检查点重新执行一次前向，即时复现出段内各层所需的中间激活，然后计算梯度。
4. 递归或嵌套：可以对段内更细粒度再应用检查点，形成递归式重算。

简而言之，就是把原来存着的东西扔掉，用的时候再算一遍。这带来了额外的计算开销（通常是前向重算一次），但显存占用量从 O(n) 降为 O(√n) 或 O(log n)，取决于分段策略。

典型实现机制

PyTorch 中的 torch.utils.checkpoint

PyTorch 提供了 torch.utils.checkpoint.checkpoint 函数，可对任意 nn.Module 或函数施加检查点。用法伪代码：

from torch.utils.checkpoint import checkpoint

class CheckpointedBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(...)
    
    def forward(self, x):
        # 第二个参数为检查点段数（通常1即整段重算）
        return checkpoint(self.layers, x, use_reentrant=False)

• use_reentrant=False（推荐）使用非重入模式，避免递归重入问题，依赖框架自动重算。
• 被检查点包裹的模块在前向时不会保存中间张量，反向时自动触发重算。

框架级支持与嵌套检查点

许多训练框架（如 DeepSpeed、Megatron-LM、Hugging Face Transformers）内置了激活检查点功能，可配置为：

• 全层检查点：每个 Transformer 层都作为检查点段。
• 选择性检查点：仅对显存占用大的模块（如注意力层）应用检查点。
• 嵌套检查点：将层分成多级段，形成更高效的重算树，进一步压榨显存（例如把注意力内部的 QKV 计算也分段）。

显存与计算的权衡分析

注意：额外的计算主要包括重新执行前向传播，不包含额外反向传播。实际训练吞吐量下降通常在 15% ～ 30% 之间，但显存节省可达 50% ～ 80%。具体数字取决于模型结构和检查点粒度。

如何在实际项目中使用

第一步：定位显存瓶颈

使用显存分析工具（如 PyTorch 的 torch.cuda.memory_summary 或 nvidia-smi）观察激活值占用是否为主要瓶颈。若峰值显存远超参数和优化器状态之和，说明激活值占用过大。

第二步：选择合适的检查点粒度

• 按层检查点：最易实现，把每个 Transformer 层包装成检查点段。适合大多数情况。
• 按模块检查点：只对注意力等显存密集型子层施加检查点，前馈网络保持原样，可在显存和速度间取得平衡。
• 嵌套检查点：仅在极限情况下使用，代码较复杂。

第三步：集成到训练代码

以 Hugging Face Transformers 为例，只需在训练参数中开启：

from transformers import Trainer, TrainingArguments

args = TrainingArguments(
    gradient_checkpointing=True,   # 激活检查点开关
    ...
)
trainer = Trainer(model=model, args=args)

或在模型代码中手动包装，确保不跟踪检查点内部梯度。

第四步：调优与验证

开启检查点后应验证：

• 训练损失曲线是否平滑下降（重算可能引入数值微小差异，通常不影响收敛）
• 显存降幅是否达到预期
• 每轮训练时间增加是否在可接受范围内

常见误区与注意事项

• 与梯度检查点的区别：激活检查点（Activation Checkpointing）有时被称为梯度检查点，容易引起混淆。它保存的是激活值用于梯度计算，而不是直接保存梯度。
• 与混合精度训练的关系：二者可叠加使用。混合精度降低张量存储字节，激活检查点减少存储张量数量，组合效果更佳。
• 随机操作（如 Dropout）重算差异：若被检查点段内包含随机操作，重算时会使用相同的随机种子以保证前向输出一致。PyTorch 在 use_reentrant=False 时会自动管理 RNG 状态。
• 不可重算的操作：部分自定义操作可能具有副作用，重算会导致错误。应避免在检查点段内包含有状态的层（如 BatchNorm 的 running mean 更新），或将其排除在检查点外。
• 显存收益并非线性：检查点自身仍需保存，分段过细反而增加检查点本身的存储开销。合理选择段长度才能最大化收益。

进阶：自适应检查点与未来方向

学术界和工业界还在研究更智能的检查点策略：

• 基于显存预算的自动分段：根据实际显存余量动态决定保留哪些层的激活。
• 激活压缩：对检查点应用有损/无损压缩（如量化），进一步降低存储。
• 重算调度优化：利用 GPU 空闲周期异步重算，隐藏计算开销。
• 反向传播与重算融合：在重算过程中即时计算梯度，避免完整保存中间张量。

这些技术正在被整合到大模型训练框架中，不断刷新“时间换空间”的极限。

总结

激活检查点是一项成熟且被广泛采用的显存优化技术，以可控的计算代价换取显存空间的大幅缩减。对于训练大模型、长序列或使用受限硬件的场景，它几乎成为必备技巧。理解其“扔掉-重算”的朴素思想，结合框架提供的简便接口，你可以轻松将其应用到自己的训练流程中，突破显存墙限制。