小内存大 batch——梯度累积

梯度累积（Gradient Accumulation）是一种训练技巧，它可以让我们在显存不那么大的 GPU 上用更大的 batch size 进行训练。

它的思路很自然：以前是在每一个 batch 上都进行参数更新；现在累积几个 batch 的梯度后，再更新参数，变相地增大了 batch size.

介绍 autograd 的那篇文章——Pytorch 中的自动求导与（动态）计算图——提到，Pytorch 的自动求导模块默认将计算好的梯度累加到各个参数的 .grad 属性上。因此在 Pytorch 中进行梯度累积很容易：

optimizer.zero_grad()
accumulation_steps = 10
for i, batch in enumerate(batches):
    # Scale the loss 
    loss = calculate_loss(batch) / accumulation_steps
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        # Reset gradients, for the next accumulated batches
        optimizer.zero_grad()

注意这一行：

loss = calculate_loss(batch) / accumulation_steps

为什么要把每一个 batch 的损失函数都除以accumulation_steps 呢？

——一般情况下，损失函数会除以 batch size 求平均。calculate_loss(batch) 求出了当下 batch size 下的平均，为了模拟更大的 batch size，还要除以 accumulation_steps。因为要模拟的 batch size 大小等于 batch size * accumulation_steps！

为了方便理解，我画了一张图。考虑两个样本的简单情况。正常情况下，batch size = 2，求出损失函数后要除以 2，然后反向传播求梯度。而在梯度累积的情况下，两个样本分别前向传播得到一个损失函数，并通过反向传播得到了梯度（accumulation steps = 2）。我们要把两个梯度加起来，用于更新参数。为了与正常情况保持一致，单样本的损失函数必须除以 2 (accumulation steps，而非 batch size)

---

假设在某个模型下，GPU 显存允许的最大 batch size 为 20。通过梯度累积的方式，我们可以实现大小为 20 * accumulation_steps 的 batch size~

先别急着走，现在考虑一种极端情况：在某个大模型下，如果 GPU 连一个样本都跑不了怎么办？——梯度检查 gradient-checkpointing！Pytorch 中的 checkpoint

---