Gradient Accumulation

Накопление градиентов — метод обучения, при котором несколько небольших батчей, называемых микробатчами, проходят прямой проход и обратное распространение ошибки, их градиенты суммируются и только затем обновляются веса модели. Такое обновление приближает один шаг оптимизации на батче, объединяющем все микробатчи, хотя в GPU одновременно хранятся активации только одного микробатча. Это позволяет обучать модель на устройстве с ограниченной памятью, используя размер батча, предусмотренный рецептом обучения, без снижения разрешения изображений или уменьшения модели. Эффективный размер батча равен размеру микробатча, умноженному на число шагов накопления и число устройств. Поэтому микробатчи по 4 изображения, накопленные за 8 шагов на 4 GPU, эквивалентны батчу из 128 изображений.

Как работает накопление градиентов#

В PyTorch и большинстве других фреймворков каждый вызов backward() добавляет новые градиенты к уже сохранённым для параметров, а не заменяет их. Накопление напрямую использует это свойство: для каждого микробатча выполняется прямой и обратный проход, оптимизатор вызывается после заданного числа микробатчей, затем градиенты очищаются и начинается следующее окно накопления:

Большинство функций потерь вычисляют среднее по примерам в батче, поэтому функцию потерь каждого микробатча делят на число шагов, чтобы суммарный градиент стал средним. При таком масштабировании четыре микробатча по 8 примеров дают тот же градиент, что и один батч из 32 примеров:

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Linear(16, 1)
x, y = torch.randn(32, 16), torch.randn(32, 1)
loss_fn = nn.MSELoss()

# Reference: one batch of 32 samples
loss_fn(model(x), y).backward()
batch_grad = model.weight.grad.clone()
model.zero_grad()

# Four micro-batches of 8: backward() adds into .grad, and dividing each loss by 4 turns the sum into a mean
steps = 4
for x_micro, y_micro in zip(x.chunk(steps), y.chunk(steps)):
    (loss_fn(model(x_micro), y_micro) / steps).backward()

print("Matches the full-batch gradient:", torch.allclose(model.weight.grad, batch_grad, atol=1e-6))
print(f"Largest difference: {(model.weight.grad - batch_grad).abs().max():.1e}")

Выводится Matches the full-batch gradient: True; наибольшая разница составляет 1.6e-07 и возникает из-за округления чисел с плавающей точкой при сложении одних и тех же значений в разном порядке.

Чем накопление отличается от большого батча#

Эквивалентность сохраняется для функций потерь, которые усредняют независимые примеры, но нарушается в трёх распространённых случаях:

• Пакетная нормализация: Вычисляет среднее и дисперсию для каждого микробатча, поэтому статистика по 4 изображениям менее точная, чем статистика по 128.
• Последовательности переменной длины: Если языковая модель усредняет функцию потерь по токенам каждого микробатча, микробатчи с малым числом токенов получают такой же вес, как и длинные.
• Обрезка градиента: Применяется один раз к накопленному градиенту непосредственно перед шагом оптимизатора; обрезка каждого микробатча меняет результат.

В октябре 2024 года Hugging Face и Unsloth устранили несоответствие в подсчёте токенов при дообучении большой языковой модели (LLM), разделив суммарную функцию потерь на общее число токенов во всём окне накопления. Накопление также не ускоряет обучение: через сеть всё равно по одному микробатчу проходят те же примеры, а небольшие микробатчи менее эффективно используют GPU. При распределённом обучении в PyTorch у DistributedDataParallel есть контекст no_sync(), который пропускает синхронизацию градиентов между устройствами для всех микробатчей, кроме последнего, сокращая объём передаваемых данных. Для большего эффективного батча обычно требуется более высокий темп обучения: исследователи Facebook обучили ResNet-50 на ImageNet за час, используя батчи по 8 192 изображения на 256 GPU, линейно увеличивая темп обучения вместе с размером батча и постепенно его разогревая.

Родственные методы#

• Контрольные точки градиентов: Повторно вычисляют активации во время обратного распространения ошибки вместо их хранения, экономя память для каждого микробатча ценой дополнительных вычислений.
• Смешанная точность: Хранит активации в 16-битном формате, освобождая место для более крупных микробатчей. Масштабатор потерь отменяет масштабирование градиентов один раз, после последнего микробатча.
• Конвейерный параллелизм: Разбивает каждый батч на микробатчи, проходящие через этапы модели на разных GPU, и накапливает их градиенты перед одним шагом оптимизатора.

Обучение Ultralytics YOLO26 по умолчанию накапливает градиенты: в режиме обучения тренер выполняет шаг оптимизатора каждые round(nbs / batch) батчей. Поэтому значение по умолчанию batch=16 и номинальный размер батча nbs=64 дают четыре батча на одно обновление, а batch=8 — восемь. Во время разогрева число постепенно увеличивается от одного, а затухание весов масштабируется на batch × accumulation / nbs, поэтому значения гиперпараметров по умолчанию, рассчитанные на батч из 64 изображений, подходят и для GPU, которая одновременно обрабатывает лишь несколько изображений.