由于没有多的 GPU 可用于张量分片（tensor sharding），又能做些什么来训练具有更大批大小（batch size）的模型呢？其中一种解决方法就是梯度累积，可以通过它来修改前面提到的训练循环。
什么是梯度积累？梯度累积是一种在训练期间虚拟增加批大小（batch size）的方法，当可用的 GPU 内存不足以容纳所需的批大小时，这非常有用。在梯度累积中，梯度是针对较小的批次计算的，并在多次迭代中累积（通常是求和或平均），而不是在每一批次之后更新模型权重。一旦累积梯度达到目标「虚拟」批大小，模型权重就会使用累积梯度进行更新。参考下面更新的 PyTorch 训练循环：
如果将 accumulation_steps 设置为 2，那么 zero_grad () 和 optimizer.step () 将只会每隔一秒调用一次。因此，使用 accumulation_steps=2 运行修改后的训练循环与将批大小（batch size）加倍具有相同的效果。例如，如果想使用 256 的批大小，但只能将 64 的批大小放入 GPU 内存中，就可以对大小为 64 的四个批执行梯度累积。（处理完所有四个批次后，将获得相当于单个批大小为 256 的累积梯度。）这样能够有效地模拟更大的批大小，而无需更大的 GPU 内存或跨不同设备的张量分片。虽然梯度累积可以帮助我们训练具有更大批量大小的模型，但它不会减少所需的总计算量。实际上，它有时会导致训练过程略慢一些，因为权重更新的执行频率较低。尽管如此，它却能帮我们解决限制问题，即批大小非常小时导致的更新频繁且混乱。例如，现在让我们运行上面的代码，批大小为 1，需要 16 个累积步骤（accumulation steps）来模拟批大小等于 16。输出如下：
... torch : 2.0.0 lightning : 2.0.0 transformers: 4.27.2 Torch CUDA available? True ... Epoch: 0001/0001 | Batch 23700/35000 | Loss: 0.0168 Epoch: 0001/0001 | Batch 24000/35000 | Loss: 0.0006 Epoch: 0001/0001 | Batch 24300/35000 | Loss: 0.0152 Epoch: 0001/0001 | Batch 24600/35000 | Loss: 0.0003 Epoch: 0001/0001 | Batch 24900/35000 | Loss: 0.0623 Epoch: 0001/0001 | Batch 25200/35000 | Loss: 0.0010 Epoch: 0001/0001 | Batch 25500/35000 | Loss: 0.0001 Epoch: 0001/0001 | Batch 25800/35000 | Loss: 0.0047 Epoch: 0001/0001 | Batch 26100/35000 | Loss: 0.0004 Epoch: 0001/0001 | Batch 26400/35000 | Loss: 0.1016 Epoch: 0001/0001 | Batch 26700/35000 | Loss: 0.0021 Epoch: 0001/0001 | Batch 27000/35000 | Loss: 0.0015 Epoch: 0001/0001 | Batch 27300/35000 | Loss: 0.0008 Epoch: 0001/0001 | Batch 27600/35000 | Loss: 0.0060 Epoch: 0001/0001 | Batch 27900/35000 | Loss: 0.0001 Epoch: 0001/0001 | Batch 28200/35000 | Loss: 0.0426 Epoch: 0001/0001 | Batch 28500/35000 | Loss: 0.0012 Epoch: 0001/0001 | Batch 28800/35000 | Loss: 0.0025 Epoch: 0001/0001 | Batch 29100/35000 | Loss: 0.0025 Epoch: 0001/0001 | Batch 29400/35000 | Loss: 0.0000 Epoch: 0001/0001 | Batch 29700/35000 | Loss: 0.0495 Epoch: 0001/0001 | Batch 30000/35000 | Loss: 0.0164 Epoch: 0001/0001 | Batch 30300/35000 | Loss: 0.0067 Epoch: 0001/0001 | Batch 30600/35000 | Loss: 0.0037 Epoch: 0001/0001 | Batch 30900/35000 | Loss: 0.0005 Epoch: 0001/0001 | Batch 31200/35000 | Loss: 0.0013 Epoch: 0001/0001 | Batch 31500/35000 | Loss: 0.0112 Epoch: 0001/0001 | Batch 31800/35000 | Loss: 0.0053 Epoch: 0001/0001 | Batch 32100/35000 | Loss: 0.0012 Epoch: 0001/0001 | Batch 32400/35000 | Loss: 0.1365 Epoch: 0001/0001 | Batch 32700/35000 | Loss: 0.0210 Epoch: 0001/0001 | Batch 33000/35000 | Loss: 0.0374 Epoch: 0001/0001 | Batch 33300/35000 | Loss: 0.0007 Epoch: 0001/0001 | Batch 33600/35000 | Loss: 0.0341 Epoch: 0001/0001 | Batch 33900/35000 | Loss: 0.0259 Epoch: 0001/0001 | Batch 34200/35000 | Loss: 0.0005 Epoch: 0001/0001 | Batch 34500/35000 | Loss: 0.4792 Epoch: 0001/0001 | Batch 34800/35000 | Loss: 0.0003 Epoch: 0001/0001 | Train acc.: 78.67% | Val acc.: 87.28% Time elapsed 51.37 min Test accuracy 87.37%
根据上面的结果，损失的波动比以前小了。此外，测试集性能提升了 10%。由于只迭代了训练集一次，因此每个训练样本只会遇到一次。训练用于 multiple epochs 的模型可以进一步提高预测性能。你可能还会注意到，这段代码的执行速度也比之前使用的批大小为 1 的代码快。如果使用梯度累积将虚拟批大小增加到 8，仍然会有相同数量的前向传播（forward passes）。然而，由于每八个 epoch 只更新一次模型，因此反向传播（backward passes）会很少，这样可更快地在一个 epoch（训练轮数）内迭代样本。结论梯度累积是一种在执行权重更新之前通过累积多个小的批梯度来模拟更大的批大小的技术。该技术在可用内存有限且内存中可容纳批大小较小的情况下提供帮助。但是，首先请思考一种你可以运行批大小的场景，这意味着可用内存大到足以容纳所需的批大小。在那种情况下，梯度累积可能不是必需的。事实上，运行更大的批大小可能更有效，因为它允许更多的并行性且能减少训练模型所需的权重更新次数。总之，梯度累积是一种实用的技术，可以用于降低小批大小干扰信息对梯度更新准确性的影响。这是迄今一种简单而有效的技术，可以让我们绕过硬件的限制。PS：可以让这个运行得更快吗？没问题。可以使用 PyTorch 2.0 中引入的 torch.compile 使其运行得更快。只需要添加一些 model = torch.compile，如下图所示：
GitHub 上提供了完整的脚本。在这种情况下，torch.compile 在不影响建模性能的情况下又减少了十分钟的训练时间：
poch: 0001/0001 | Batch 26400/35000 | Loss: 0.0320 Epoch: 0001/0001 | Batch 26700/35000 | Loss: 0.0010 Epoch: 0001/0001 | Batch 27000/35000 | Loss: 0.0006 Epoch: 0001/0001 | Batch 27300/35000 | Loss: 0.0015 Epoch: 0001/0001 | Batch 27600/35000 | Loss: 0.0157 Epoch: 0001/0001 | Batch 27900/35000 | Loss: 0.0015 Epoch: 0001/0001 | Batch 28200/35000 | Loss: 0.0540 Epoch: 0001/0001 | Batch 28500/35000 | Loss: 0.0035 Epoch: 0001/0001 | Batch 28800/35000 | Loss: 0.0016 Epoch: 0001/0001 | Batch 29100/35000 | Loss: 0.0015 Epoch: 0001/0001 | Batch 29400/35000 | Loss: 0.0008 Epoch: 0001/0001 | Batch 29700/35000 | Loss: 0.0877 Epoch: 0001/0001 | Batch 30000/35000 | Loss: 0.0232 Epoch: 0001/0001 | Batch 30300/35000 | Loss: 0.0014 Epoch: 0001/0001 | Batch 30600/35000 | Loss: 0.0032 Epoch: 0001/0001 | Batch 30900/35000 | Loss: 0.0004 Epoch: 0001/0001 | Batch 31200/35000 | Loss: 0.0062 Epoch: 0001/0001 | Batch 31500/35000 | Loss: 0.0032 Epoch: 0001/0001 | Batch 31800/35000 | Loss: 0.0066 Epoch: 0001/0001 | Batch 32100/35000 | Loss: 0.0017 Epoch: 0001/0001 | Batch 32400/35000 | Loss: 0.1485 Epoch: 0001/0001 | Batch 32700/35000 | Loss: 0.0324 Epoch: 0001/0001 | Batch 33000/35000 | Loss: 0.0155 Epoch: 0001/0001 | Batch 33300/35000 | Loss: 0.0007 Epoch: 0001/0001 | Batch 33600/35000 | Loss: 0.0049 Epoch: 0001/0001 | Batch 33900/35000 | Loss: 0.1170 Epoch: 0001/0001 | Batch 34200/35000 | Loss: 0.0002 Epoch: 0001/0001 | Batch 34500/35000 | Loss: 0.4201 Epoch: 0001/0001 | Batch 34800/35000 | Loss: 0.0018 Epoch: 0001/0001 | Train acc.: 78.39% | Val acc.: 86.84% Time elapsed 43.33 min Test accuracy 87.91%
请注意，与之前相比准确率略有提高很可能是由于随机性。
原文链接：https://lightning.ai/pages/blog/gradient-accumulation/
