Как обучаться на большом батче при ограниченной памяти?

Что проверяют

Различаете ли вы размер батча и размер датасета и знаете ли приёмы, которые меняют память на вычисления. Следом обычно спрашивают, эквивалентно ли накопление градиентов настоящему большому батчу и в каких случаях нет.

Как отвечать

Основной приём это накопление градиентов. Большой батч делится на k микробатчей, для каждого вызывается backward(), градиенты складываются в .grad, а optimizer.step() и zero_grad() вызываются раз в k итераций. Лосс каждого микробатча делят на k, иначе градиент выйдет в k раз больше. В DDP промежуточные шаги оборачивают в model.no_sync(), чтобы не синхронизировать градиенты между GPU на каждом микробатче.

Дальше то, что уменьшает память на один пример:

• смешанная точность: torch.autocast с bf16 или fp16; для fp16 нужен GradScaler, bf16 обычно обходится без него;
• чекпоинтинг активаций: torch.utils.checkpoint не хранит промежуточные активации, а пересчитывает их на обратном проходе. Плата примерно один лишний прямой проход, экономия заметная;
• состояние оптимизатора: Adam хранит два дополнительных тензора на каждый параметр. Помогают 8-битные оптимизаторы, шардирование через FSDP или DeepSpeed ZeRO, выгрузка состояния на CPU;
• больше GPU: глобальный батч равен размеру на устройство, умноженному на число устройств и шагов накопления.

При смене эффективного батча пересматривают learning rate: линейное масштабирование с разогревом, а для очень больших батчей оптимизаторы вроде LAMB.

Где спотыкаются

Отвечают про загрузку датасета по частям, генераторы и PCA. Это ответ на другой вопрос: данные и так идут батчами, а память кончается на активациях, градиентах и состоянии оптимизатора.

Забывают про BatchNorm. Статистики считаются по микробатчу, а не по накопленному батчу, поэтому при крошечных микробатчах лучше GroupNorm, LayerNorm или синхронизированная нормализация.

Считают накопление полным эквивалентом. В контрастивных лоссах вроде InfoNCE негативы берутся только внутри микробатча, и накопление их число не увеличивает. Для этого есть отдельные приёмы, например GradCache.