# 梯度累积

较大的批次会产生较大的激活值，从而耗尽 GPU 内存。梯度累积允许你通过跨多个小批次分散梯度计算，来使用更大的有效批次大小进行训练。

梯度会在 *n* 个 mini-batch 上累积，之后优化器才会更新权重。例如，当每个设备的 batch size 为 8 且累积步数为 4 时，有效 batch size 为 32。

```text
Step 1: mini-batch 1 → forward → backward → grads = G₁
Step 2: mini-batch 2 → forward → backward → grads = G₁ + G₂
Step 3: mini-batch 3 → forward → backward → grads = G₁ + G₂ + G₃
Step 4: mini-batch 4 → forward → backward → grads = G₁ + G₂ + G₃ + G₄
        → optimizer.step()  ← same update as if batch_size × 4
        → zero_grad()
```

仅当更大的批次无法放入内存时，才使用梯度累积。与使用真正的大批次进行训练相比，它并不能提高吞吐量。

在 `per_device_train_batch_size` 的基础上，累积 `gradient_accumulation_steps` 步的梯度。

```py
from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
)
```

## 损失缩放

对于[自定义损失函数](./trainer_recipes#custom-loss-function)，请包含`num_items_in_batch`，以便[Trainer](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer)将损失除以所有 mini-batch 中预测目标的总数。这样会按 token 数量进行归一化，而不是按`gradient_accumulation_steps`的固定步数。否则，[Trainer](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer)会将损失除以`gradient_accumulation_steps`。

```py
import torch.nn.functional as F

def compute_loss(outputs, labels, num_items_in_batch=None):
    logits = outputs["logits"]
    loss = F.cross_entropy(logits, labels, reduction="sum")
    return loss / num_items_in_batch
```

对于因果语言模型，`num_items_in_batch` 统计的是经过*移位*后的标签。损失函数会对标签进行移位，使得位置 `i` 处的预测目标为位置 `i + 1` 处的 token，这导致每个序列的位置 0 没有对应的目标。[Trainer](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer) 会排除这些位置，并基于 `labels[..., 1:]` 进行统计，因此分母与损失函数实际使用的预测目标数量保持一致。当数据整理器（data collator）直接提供 `shift_labels` 时（例如无填充的整理器），[Trainer](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer) 会改为基于该张量进行统计。其他类型的损失函数，如掩码语言模型和分类任务，则统计完整的标签张量。

如果你子类化 [Trainer](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer) 并重写 [compute_loss()](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer.compute_loss)，请设置 `loss_is_scaled_for_ga` 以说明你的损失是如何归一化的：对于按批次求平均的情况，请设为 `False`，[Trainer](/docs/transformers/main/zh/main_classes/trainer#transformers.Trainer) 随后会将其除以 `gradient_accumulation_steps`；如果你已经将其除以 `num_items_in_batch`，则设为 `True`。

```py
class MyTrainer(Trainer):
    loss_is_scaled_for_ga = True

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        ...
        return loss_sum / num_items_in_batch
```

## 后续步骤

- 阅读[GPU 内存使用](./model_memory_anatomy)文档，以了解训练期间 GPU 内存使用量的驱动因素。
- 请参阅 [梯度检查点](./grad_checkpointing) 指南，了解如何通过重新计算激活值而非缓存激活值来减少激活内存。
- 请参阅[混合精度训练](./mixed_precision_training)指南，了解如何使用较低精度的数据类型来减少内存占用并加快训练速度。
- 阅读 [Gradient Accumulation Fix](https://unsloth.ai/blog/gradient) 博客文章，了解梯度累积的计算方式。