activation-checkpointing

激活检查点：用反向重算换显存，训练长上下文和大模型的基础手段

补充说明

激活检查点（activation checkpointing，也叫 gradient checkpointing）在前向时不保存全部中间激活，只保留少量检查点；反向传播时重新计算中间结果。它把显存换成了额外计算，通常能显著降低激活占用，是长上下文、大 batch 和多卡训练的基础优化。

PyTorch 原生用 torch.utils.checkpoint，Hugging Face 模型提供 gradient_checkpointing_enable()，FSDP/DeepSpeed/ColossalAI 也各有集成方式。它只省激活显存，不省参数、优化器状态和 KV Cache；推理阶段通常不需要。

常用示例

# Hugging Face Transformers
model.gradient_checkpointing_enable()
model.config.use_cache = False  # 训练时关闭 KV Cache

import torch
from torch.utils.checkpoint import checkpoint

def block(x):
    return model_block(x)

# 非重入模式，支持更多场景（推荐）
y = checkpoint(block, x, use_reentrant=False)

# FSDP：按 Transformer block 开启激活检查点
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper, CheckpointImpl, apply_activation_checkpointing,
)
from transformers.models.llama.modeling_llama import LlamaDecoderLayer

apply_activation_checkpointing(
    model,
    checkpoint_wrapper_fn=lambda m: checkpoint_wrapper(
        m, checkpoint_impl=CheckpointImpl.NO_REENTRANT
    ),
    check_fn=lambda m: isinstance(m, LlamaDecoderLayer),
)

# DeepSpeed 配置片段
{
  "activation_checkpointing": {
    "partition_activations": True,
    "contiguous_memory_optimization": True
  }
}

易错点与排错

• 以为能省所有显存：它只省激活，不省参数、优化器状态和 KV Cache；显存仍不够时要结合 FSDP/ZeRO/量化。
• 训练时忘记关闭 KV Cache：Hugging Face 模型训练时 use_cache=True 会额外占用显存；设为 False。
• 重入模式与非重入模式混用：use_reentrant=True 对某些控制流/RNG 处理不同；新代码优先 use_reentrant=False。
• Dropout/RNG 不一致：重算时随机数状态要一致，否则前后向不一致；非重入模式通常处理更好，仍要用固定种子验证。
• 与 FSDP 包装顺序错误：先包装模型再应用 checkpoint wrapper，或按框架文档顺序；顺序错会漏掉层或报错。
• 全模型每个算子都 checkpoint：过度使用会显著变慢；按 Transformer block 粒度开启，平衡显存和计算。
• 推理阶段也开启：推理没有反向，checkpoint 没有收益，反而增加开销；部署时关闭。
• 长上下文仍然 OOM：检查点降低激活，但注意力矩阵和 KV 仍可能很大；结合 FlashAttention、序列并行或分块。

相关命令

• fsdp — 参数/梯度/优化器分片。
• deepspeed — ZeRO 与激活检查点集成。
• torch-compile — 编译优化与重算。
• liger-kernel — 融合算子降低激活。
• distributed-training-troubleshooting — OOM 与挂起排障。

参考链接

• PyTorch checkpoint 文档
• Hugging Face gradient checkpointing
• FSDP 激活检查点

读完自测

合上页面，先回答，再展开对照。

1. 激活检查点用什么换什么？它不省哪些显存？
2. 训练时为什么要关闭 use_cache？
3. use_reentrant=True 和 False 应该怎么选？