Накопление градиентов

Крупные батчи порождают большие активации, которые исчерпывают память GPU. Накопление градиентов позволяет обучать модель с большим эффективным размером батча, распределяя вычисление градиентов по нескольким мини-батчам.

Градиенты накапливаются в течение n мини-батчей, прежде чем оптимизатор обновляет веса. Например, при размере батча на устройство 8 и 4 шагах накопления эффективный размер батча составляет 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()

Используйте накопление градиентов только в том случае, если более крупная партия не помещается в память. Это не повышает пропускную способность по сравнению с обучением на действительно крупной партии.

Накапливайте градиенты в течение gradient_accumulation_steps для per_device_train_batch_size.

from transformers import TrainingArguments

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

Масштабирование потерь

Для собственной функции потерь включите num_items_in_batch, чтобы Trainer делил потерю на количество целевых значений предсказания во всех мини-батчах. Это нормализует значение по токенам, а не по фиксированному числу шагов с 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

Для моделей каузального LM параметр num_items_in_batch учитывает сдвинутые метки. Потери сдвигают метки так, чтобы предсказание в позиции i соответствовало токену в позиции i + 1, из-за чего позиция 0 каждой последовательности не имеет целевого значения. Trainer исключает эти позиции и считает по labels[..., 1:], поэтому знаменатель совпадает с количеством целевых значений для предсказаний, используемых в потере. Если коллабор данных напрямую предоставляет shift_labels, например, коллабор без паддинга, Trainer считает по этому тензору. Другие типы потерь, такие как masked LM и классификация, учитывают весь тензор меток.

Если вы создаёте подкласс Trainer и переопределяете compute_loss(), установите loss_is_scaled_for_ga, чтобы указать, как нормализована ваша функция потерь: False для среднего по батчу, которое Trainer затем делит на gradient_accumulation_steps, или True, если вы уже разделили её на num_items_in_batch.

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 checkpointing, чтобы узнать, как уменьшить память, занимаемую активациями, за счёт их пересчёта вместо кэширования.
• См. руководство Обучение с смешанной точностью, чтобы узнать, как использовать типы данных с пониженной точностью для сокращения потребления памяти и ускорения обучения.
• Ознакомьтесь с публикацией Gradient Accumulation Fix, чтобы узнать, как вычисляется накопление градиентов.