人工智能：梯度累计 GRADIENT_ACCUMULATION_STEPS

Posted on Edited on In Waline:

GRADIENT_ACCUMULATION_STEPS（梯度累积步数）是一种在不增加显存（VRAM）消耗的情况下，模拟出更大批量（Batch Size）训练效果的关键技术。

技术原理详解

在标准的神经网络训练中，一个训练步骤（step）包含以下过程：

  1. 前向传播: 输入一个批次的数据，模型计算出预测结果和损失（loss）。
  2. 后向传播: 根据损失计算出每个参数的梯度（gradient），即参数应该调整的方向。
  3. 优化器步骤: 优化器（如 AdamW）根据梯度来更新模型的权重。
  4. 清空梯度: 为下一次计算做准备。

而当使用梯度累积时，这个过程发生了变化：

对于第 1 到 N-1 步 (N = GRADIENT_ACCUMULATION_STEPS):

  1. 前向传播: 正常执行。
  2. 后向传播: 正常执行，计算出当前“微批次 (micro-batch)”的梯度。
  3. 关键区别: 跳过“优化器步骤”。我们不立即更新权重，而是让梯度累积在模型的 .grad 属性中。同时，我们也不清空梯度。

对于第 N 步:

  1. 前向传播: 正常执行。
  2. 后向传播: 正常执行，计算出第 N 个微批次的梯度，并将其加到之前已经累积的梯度上。
  3. 执行“优化器步骤”: 优化器使用累积了 N 步的总梯度来对模型权重进行一次更新。
  4. 清空梯度: 更新完成后，将梯度清零，准备开始下一个累积周期。

这样一来，虽然硬件上每次只处理了一个小批量，但参数的更新却是基于多个小批量梯度的总和，从而在数学效果上模拟了一个大批量的训练。

核心公式

有效批量大小 (Effective Batch Size) = PER_DEVICE_TRAIN_BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS * GPU数量

例如，在单张 A100 上，即使你的 PER_DEVICE_TRAIN_BATCH_SIZE 因为显存限制只能设为 2，但只要将 GRADIENT_ACCUMULATION_STEPS 设为 32，你就能达到 2 * 32 = 64 的有效批量大小。

优缺点分析

优点 (Pros)

  1. 突破显存限制: 这是其最核心的优点。它允许我们在有限的硬件资源上训练需要大批量才能稳定收敛的超大型模型。
  2. 训练过程更稳定: 更大的有效批量意味着每次权重更新所依据的梯度方向更准确、噪声更小，这使得整个训练过程更加平稳，允许你使用更高的学习率。
  3. 可能提升模型性能: 稳定的训练过程和更准确的梯度方向，有助于模型找到更好的局部最小值，从而可能提升最终的泛化能力。

缺点 (Cons)

  1. 训练时间变长 (Wall-clock Time): 这是最大的代价。虽然模拟了大批量，但计算过程仍然是串行的。权重更新的频率降低了，完成一个 epoch 所需的实际时间会相应增加。例如，累积32步才更新一次，意味着权重更新的频率是原来的 1/32，完成一个 epoch 的总时间大约会是原来的数倍（具体取决于数据加载等其他开销）。
  2. 对特定层（如 BatchNorm）的影响: 对于包含批归一化（Batch Normalization）层的模型，梯度累积可能会带来问题。因为 BatchNorm 是在每个微批次上计算均值和方差的，而不是在整个“有效批量”上。这可能导致训练和推理时的统计数据不匹配。不过，现代大型语言模型（LLM）大多使用 Layer Normalization 或 RMS Normalization，它们不受批量大小的影响，因此这个问题在 LLM 训练中基本不存在。

总结: GRADIENT_ACCUMULATION_STEPS 是一种典型的 用时间换显存的策略。它是大规模模型训练中不可或缺的利器，使得个人或小型机构在有限的硬件条件下训练强大的模型成为可能。