Feat: Activation Checkpointing

[Chamberlain0w0]

背景

本 PR 为 InfiniTrain 增加重计算能力，通过在反向传播时重新执行部分前向计算，减少训练过程中需要保留的中间激活，以计算开销换取显存占用下降。

整体配置与 Megatron-Core 的 Transformer activation recompute 语义保持一致。

设计

通用 Checkpoint 接口

新增 utils::checkpoint::Checkpoint()，对标 torch.utils.checkpoint()。同时支持：

• Reentrant checkpoint
• Non-reentrant checkpoint，默认使用该模式

Non-reentrant 模式通过 Saved Tensor Hooks 将前向过程中保存的中间 Tensor 替换为占位对象。反向传播真正读取这些 Tensor 时，再按需重新执行 checkpoint 区域的前向计算。

该实现包含以下机制：

• Early stop：需要的中间结果全部生成后，提前终止重计算
• Determinism check：校验重计算结果的 shape、dtype 和 device 是否与原始前向一致
• Autocast context 恢复：重计算沿用原始前向的 autocast 配置，保证 BF16 等混合精度场景的一致性
• 输入 Tensor detach 后重建 requires_grad 状态，避免错误连接原始计算图
• 重计算结束后及时释放 Tensor 和引用，降低峰值显存并避免引用环

由于当前 autograd engine 在构图时会直接修改依赖计数，non-reentrant 重计算在 no_grad 下执行，并新增 PropagateRequiresGradGuard 传递 requires_grad 和 needs_input_grad 元数据，避免临时重计算图污染原始反向图。

Transformer 重计算策略

TransformerConfig 新增以下配置：

• recompute_granularity
  • none：关闭重计算
  • full：以完整 Transformer Layer 为重计算单位
  • selective：预留配置，当前尚未实现
• recompute_method
  • none：每个 Transformer Layer 独立 checkpoint
  • uniform：每连续 recompute_num_layers 层组成一个 checkpoint 区域
  • block：仅 checkpoint 当前本地 chunk/stage 的前 recompute_num_layers 层
• recompute_num_layers：控制 uniform 分组大小或 block 覆盖层数

未启用重计算或处于 no_grad 场景时，Transformer 仍直接执行原有前向路径。

主要修改

• 新增通用 activation checkpoint API，以及 reentrant/non-reentrant 实现
• 扩展 autograd grad mode，使 non-reentrant 重计算可以传递梯度相关元数据而不构建临时计算图
• 在 Transformer Layer 执行路径中接入 full activation recompute
• 为 GPT-2 和 LLaMA 3 示例新增命令行参数：
  • --activation_recompute
  • --recompute_granularity
  • --recompute_method
  • --recompute_num_layers
• LLMC checkpoint loader 支持将运行时 recompute 配置应用到加载后的模型
• 统一 Module 调用入口，确保前向 hooks 等 Module 调用协议能够正常执行
• 将 MNIST 示例中的直接 Forward() 调用调整为 operator()

当前限制

• selective granularity 尚未实现，目前仅支持完整 Transformer Layer 重计算
• RNG state 保存与恢复尚未实现，因此 checkpoint 区域暂不支持依赖随机状态一致性的 Dropout 或其他随机算子
• Reentrant checkpoint 当前仅支持单输出 forward

[chen2021673]

之后扩展 recompute 其他功能（如selective）的话就不适合全放在transformer.cc里了，之后可以拆分activation_recompute.cc。本PR可以不做修改

[Chamberlain0w0]

确实，这块确实没太想清楚，主要是 megatron 的实现里面也把很多重计算逻辑融在了 transformer 模型层，之后可以再讨论下

[Chamberlain0w0]

本质上是需要一个 2*2=4 种情况的精细控制（with_grad/no_grad x 带引用计数建图/仅传递 requires_grad）。

也可以不单独实现一个 guard，而是给原先的 GradGuard 重载一个带参数 (比如 bool record_context = false) 的构造方式，让 GradGuard 本身能够覆盖这四种情况。

[Chamberlain0w0]

只是用异常的语法和控制流机制来实现 early stop 的语义，本质上并不是真的抛一个会让程序 abort 的异常。torch 也采取的是一样的实现方式。

这块暂时没有想到其他更优雅的处理方式。