Gradient Checkpointing

梯度检查点是一种节省内存的训练技术，它仅存储前向传播中选定的中间激活值，并在反向传播期间重新计算其余部分。它也被称为激活检查点，通过增加额外的计算量来换取更低的峰值内存使用量。尽管名字叫梯度检查点，但该技术检查点化的是激活值，而不是参数梯度或模型文件，这在激活张量导致神经网络无法放入可用GPU内存时显得尤为珍贵。

梯度检查点的工作原理#

在标准前向传播期间，神经网络会计算被称为激活的中间张量。正如PyTorch autograd机制中所述，自动微分系统会保留稍后计算梯度所需的激活值。深层网络、大批量、高分辨率图像和长输入序列会使这些保存的张量消耗大量的内存。

梯度检查点将网络划分为若干段：

1. 前向传播会存储选定段的输入或边界激活值。
2. 那些段内的其他中间激活值将被丢弃。
3. 在反向传播期间，每个检查点化的段会再次向前运行以重建丢失的值。
4. 重建的激活值会立即用于计算梯度。

PyTorch激活检查点API通过 torch.utils.checkpoint 暴露出这一行为。同类概念还包括JAX梯度检查点与重新计算以及TensorFlow带区检查点。

检查点的位置决定了权衡取舍。检查点化更多的区域通常能节省更多内存，但会重复更多的运算。在激活密集的块周围进行选择性放置，可以比重新计算整个网络提供更好的平衡。

权衡与相关技术#

梯度检查点通常保持模型参数、梯度和优化器状态不变。它的主要目标是激活内存，并且由于某些前向计算运行两次，训练会变慢。具体结果取决于架构、检查点边界、批次形状和硬件。

它与几种相关技术有所不同：

• **梯度累积**在更新模型权重之前处理多个微批次，从而在无需同时加载每个样本的情况下创建更大的有效批次。而梯度检查点则是减少为每个微批次保留的激活值。
• **混合精度**对选定操作使用较低精度的的数据类型。PyTorch自动混合精度工作流可以减少内存并加速兼容的操作，而检查点则刻意增加了计算量。
• 批次大小缩减通过一起处理较少的样本来降低内存。检查点可能允许更大的批次或输入分辨率保持可行。
• 训练检查点保存权重和优化器状态以便恢复或后续推理。PyTorch模型检查点指南描述了这种持久化机制，它与激活值的重新计算无关。

检查点代码在其原始和重新计算的前向执行之间应具有功能一致性。检查点区域内的可变状态、设备传输或不受控制的随机性可能会导致错误或不正确的梯度。在解码器风格的Transformer中，在进行检查点训练期间可能还需要禁用键值缓存，因为缓存的推理状态可能会与重建前向图产生冲突。

PyTorch示例#

以下示例在单个训练步骤中对内存密集型块进行了检查点化处理：

import torch
from torch import nn
from torch.utils.checkpoint import checkpoint

torch.manual_seed(0)

block = nn.Sequential(
    nn.Linear(1024, 4096),
    nn.ReLU(),
    nn.Linear(4096, 1024),
)
optimizer = torch.optim.AdamW(block.parameters())
inputs = torch.randn(8, 1024, requires_grad=True)
targets = torch.zeros_like(inputs)

optimizer.zero_grad(set_to_none=True)
outputs = checkpoint(block, inputs, use_reentrant=False)
loss = nn.functional.mse_loss(outputs, targets)
loss.backward()
optimizer.step()

仅保留该块的输入和所需的边界信息；其内部激活值在 loss.backward() 期间被重建。实际项目应使用诸如PyTorch峰值分配GPU内存之类的度量指标来比较启用和不启用检查点时的运行情况。

实际应用#

• **高分辨率计算机视觉：**医学分割、航空检测和工业检查可能会对大图像进行训练，其特征图消耗的内存比模型权重还要多。对选定的主干阶段进行检查点化可以保持图像分辨率，或者允许每个批次处理更多样本，而不是激进地对输入进行降采样。

• **长序列Transformer训练：**随着序列长度和层数的增加，激活存储会迅速增长。重新计算Transformer块可以使更长的上下文或更大的微批次适应相同的加速器。NVIDIA激活重新计算指南阐述了Transformer层的完整和选择性重新计算。

实用指南#

当性能分析显示激活值（而非参数或优化器状态）占据内存主导时，请使用梯度检查点。从大的重复块开始，并在扩大覆盖范围之前对峰值内存、迭代时间和验证行为进行基准测试。

对于文档化的Ultralytics YOLO训练配置，一线内存控制包括自动混合精度、图像大小和物理批次大小。Ultralytics AutoBatch参考解释了基于可用GPU内存的自动批次选择。当本地硬件依然不足时，Ultralytics平台云训练为托管训练运行提供可配置的云GPU。当自定义PyTorch架构需要更精细的激活内存管理时，梯度检查点可以作为这些控制措施的补充。