AI 棱镜站

主题

梯度检查点节省显存 ​

梯度检查点（Gradient Checkpointing）通过选择性地丢弃前向传播的中间激活值，在反向传播时重新计算，以额外计算开销换取显存节省。

梯度检查点

原理 ​

标准训练中，前向传播保存所有中间激活值供反向传播使用：

显存占用 = O(层数 × 批次 × 序列长度 × 隐藏维度)

梯度检查点只保存部分层的激活值（检查点），其余层在反向传播时重新计算：

显存占用 = O(√层数 × 批次 × 序列长度 × 隐藏维度)
计算开销 = 额外 33% 前向计算

PyTorch 实现 ​

python
from torch.utils.checkpoint import checkpoint

class CheckpointedTransformerLayer(nn.Module):
    def __init__(self, layer):
        super().__init__()
        self.layer = layer

    def forward(self, x):
        # 使用梯度检查点包装
        return checkpoint(self.layer, x, use_reentrant=False)

# 或使用 PyTorch 内置方法
model.gradient_checkpointing_enable()

Hugging Face Transformers 配置 ​

python
from transformers import LlamaForCausalLM, LlamaConfig

# 启用梯度检查点
model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    torch_dtype=torch.bfloat16,
)
model.gradient_checkpointing_enable(
    gradient_checkpointing_kwargs={"use_reentrant": False}
)

选择性检查点 ​

python
# 只对 Attention 层使用检查点
def selective_checkpoint_fn(module):
    return isinstance(module, (LlamaAttention, GPTNeoXAttention))

from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing,
    checkpoint_wrapper,
)

apply_activation_checkpointing(
    model,
    checkpoint_wrapper_fn=checkpoint_wrapper,
    check_fn=selective_checkpoint_fn,
)

显存节省效果 ​

模型无检查点全部检查点选择性检查点
7B (BS=1, Seq=2048)16GB8GB10GB
13B (BS=1, Seq=2048)28GB14GB18GB
70B (BS=1, Seq=2048)140GB70GB90GB

检查点策略

  • 全部检查点：最大显存节省，额外 33% 计算开销
  • 选择性检查点：仅对 Attention 检查点，MLP 保留激活值，开销约 15%
  • 推荐：对 Attention 使用检查点，MLP 不使用，兼顾显存和速度

训练速度影响

梯度检查点增加约 20-33% 的训练时间。在显存充足时不应启用。建议仅在 OOM 或需要增大 batch size 时使用。

相关资源 ​

最近更新