数学原理参考：

梯度检查点技术（Gradient Checkpointing）详细介绍：中英双语-CSDN博客

视频讲解参考：

用梯度检查点来节省显存 gradient checkpointing_哔哩哔哩_bilibili

Gradient Checkpointing（梯度检查点

Gradient Checkpointing 是一种用于优化深度学习模型训练的技术，旨在减少训练过程中显存的占用。在深度神经网络训练中，通常需要存储每一层的激活值以用于反向传播计算梯度。然而，对于层数较多或参数量较大的模型，这些激活值会占用大量显存。

Gradient Checkpointing 的核心思想是在前向传播时选择性地保存部分激活值（称为检查点），而丢弃其他激活值。在反向传播时，如果需要这些被丢弃的激活值，则重新计算它们。通过这种方式，显存使用量可以从 O(L) 降低到 O(K)，其中 L 是网络层数，K 是选择的检查点层数。

工作原理

1. 选择检查点：在前向传播时，选择某些层作为检查点，保存这些层的激活值。

2. 丢弃激活值：对于未被选为检查点的层，丢弃其激活值。

3. 反向传播时重新计算：在反向传播时，如果需要被丢弃的激活值，则通过重新计算它们来获取，从而计算梯度。

a1和a3被丢弃，反向传播时，如果需要被丢弃的激活值，则需要重新计算

a1 = x * w1,

a3 = a2 * w3

优点与缺点

优点：

• 显著减少显存占用，使训练更大规模的模型成为可能。

• 在显存受限的环境中，可以提高训练效率。

• 允许使用更大的批量大小，从而加速训练。

缺点：

• 增加了计算开销，因为需要在反向传播时重新计算激活值。

• 实现复杂度增加，需要修改代码来管理检查点。

• 可能导致训练时间延长。

实现方法

在 PyTorch 中，可以通过 torch.utils.checkpoint 模块实现 Gradient Checkpointing。例如：

import torch
import torch.nn as nn
import torch.utils.checkpoint as checkpointclass SimpleModel(nn.Module):def __init__(self):super(SimpleModel, self).__init__()self.layer1 = nn.Linear(256, 256)self.layer2 = nn.Linear(256, 256)self.layer3 = nn.Linear(256, 10)def forward(self, x):x = checkpoint.checkpoint(self.layer1, x)  # 应用梯度检查点x = checkpoint.checkpoint(self.layer2, x)x = self.layer3(x)  # 最后一层不需要检查点return x

在 DeepSpeed 中，可以通过配置文件启用 Gradient Checkpointing：

{"train_batch_size": 16,"gradient_accumulation_steps": 4,"zero_optimization": {"stage": 2,"contiguous_gradients": true},"gradient_checkpointing": true
}

应用场景

Gradient Checkpointing 广泛应用于以下场景：

• 训练大规模深度学习模型，如 7B 或 10B 参数的模型。

• 在 GPU 显存有限的环境中优化训练。

• 提高训练效率，同时减少硬件成本。

通过合理使用 Gradient Checkpointing，可以在有限的硬件资源下训练更大规模的模型，同时平衡显存和计算开销。