梯度累积

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

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

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 步的梯度。

from transformers import TrainingArguments

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

损失缩放

对于自定义损失函数，请包含num_items_in_batch，以便Trainer将损失除以所有 mini-batch 中预测目标的总数。这样会按 token 数量进行归一化，而不是按gradient_accumulation_steps的固定步数。否则，Trainer会将损失除以gradient_accumulation_steps。

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 会排除这些位置，并基于 labels[..., 1:] 进行统计，因此分母与损失函数实际使用的预测目标数量保持一致。当数据整理器（data collator）直接提供 shift_labels 时（例如无填充的整理器），Trainer 会改为基于该张量进行统计。其他类型的损失函数，如掩码语言模型和分类任务，则统计完整的标签张量。

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

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 内存使用文档，以了解训练期间 GPU 内存使用量的驱动因素。
• 请参阅 梯度检查点 指南，了解如何通过重新计算激活值而非缓存激活值来减少激活内存。
• 请参阅混合精度训练指南，了解如何使用较低精度的数据类型来减少内存占用并加快训练速度。
• 阅读 Gradient Accumulation Fix 博客文章，了解梯度累积的计算方式。