智体通鉴 · AgentGuide

外观

GPU显存优化策略完全指南 ​

内容整理自学习笔记，仅供面试备考参考；不构成录用、培训或考试承诺。

1. 什么是梯度累积（Gradient Accumulation）？ ​

梯度累积是一种在显存受限情况下模拟大batch训练的技术。其核心思想是将多个小batch的梯度累积起来，达到指定步数后一次性更新参数。

传统梯度更新方式： 每个batch计算损失后立即更新参数：

python
for (inputs, labels) in data_loader:
    inputs = inputs.to(device)
    labels = labels.to(device)

    with torch.set_grad_enabled(True):
        preds = model(inputs)
        loss = criterion(preds, labels)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

梯度累积方式： 每个batch只计算梯度并累积，达到 gradient_accumulation_steps 次后才更新参数：

python
gradient_accumulation_steps = 4

for batch_idx, (inputs, labels) in enumerate(data_loader):
    inputs = inputs.to(device)
    labels = labels.to(device)

    with torch.set_grad_enabled(True):
        preds = model(inputs)
        loss = criterion(preds, labels)
        loss /= gradient_accumulation_steps
        loss.backward()

    if ((batch_idx + 1) % gradient_accumulation_steps == 0) or \
       ((batch_idx + 1) == len(data_loader)):
        optimizer.step()
        optimizer.zero_grad()

注意： loss需要除以累积步数，以保证梯度均值与真实大batch一致。

梯度累积的优势 ​

优势说明
内存效率累积过程不占用额外内存，允许在显存有限时使用更大的等效batch
训练稳定性更大的等效batch减少梯度方差，提供更稳定的梯度信号
更快收敛大batch数据包含更丰富的信息，有助于更快收敛到较好的模型状态
参数更新频率控制通过设置累积步数，灵活控制更新频率

注意事项： 较大的累积步数可能导致更新频率过低、训练速度下降；累积梯度也可能影响动量和学习率等参数的计算。

2. 什么是梯度检查点（Gradient Checkpointing）？ ​

梯度检查点是一种用时间换空间的显存优化技术。它通过在计算图中插入检查点，仅保存部分中间结果，反向传播时重新计算被丢弃的中间结果，从而大幅减少显存占用。

背景 ​

深度神经网络的反向传播需要存储大量中间梯度，随着模型规模增大，中间结果的显存占用急剧上升。

工作原理 ​

  1. 前向传播时，只保存部分关键中间结果（检查点），其余中间结果算完即丢弃
  2. 反向传播时，从检查点重新进行前向计算，恢复所需的中间结果
  3. 以额外计算开销换取大幅显存节省

梯度检查点的优势 ​

优势说明
减少显存压力大幅降低反向传播所需的显存，支持有限硬件训练更大模型
支持大batch训练允许在大batch训练中有效使用显存
控制显存开销无需昂贵硬件即可进行大规模实验
python
# PyTorch中使用gradient checkpointing
from torch.utils.checkpoint import checkpoint

# 在模型forward中使用
def forward(self, x):
    return checkpoint(self._forward, x)

注意事项： 梯度检查点虽然节省显存，但会引入额外的计算开销（需要重算前向过程），实际使用时需权衡计算时间和显存占用。

3. 梯度累积与梯度检查点对比 ​

对比维度梯度累积梯度检查点
优化目标等效batch大小中间结果显存占用
核心策略多batch累积再更新丢弃中间结果、反向时重算
时间开销无额外计算开销需额外前向重计算
显存收益不增加显存但模拟大batch大幅减少反向传播显存
适用场景batch受限、显存够但batch小模型太大、中间结果占显存过多

实践建议： 两者可以同时使用——梯度累积扩大等效batch，梯度检查点降低单batch的显存占用，共同实现在有限显存下训练大模型。