Накопление градиентов: эффективные батчи с правильной нормировкой

Определение

При накоплении градиентов несколько микробатчей обрабатываются до обновления параметров модели. Их градиенты участвуют в одном шаге оптимизатора, позволяя увеличить эффективный батч, когда физический батч не помещается в память устройства. Вызовы backward в PyTorch прибавляют результаты к существующим градиентам параметров, пока те не будут очищены. [1][2]

Принцип работы

Для микробатчей одинакового размера и средней потери по одинаково взвешенным примерам делите каждую потерю на число микробатчей в окне накопления. Тогда сумма градиентов совпадает с градиентом общей средней потери. При неодинаковых размерах нужны веса по фактическому числу участвующих примеров. Для потерь по токенам считайте допустимые незамаскированные токены; веса классов и другие способы агрегации могут требовать иного знаменателя. Бездумное усреднение средних микробатчей меняет целевую функцию. Эти правила следуют из дифференцирования нужного взвешенного среднего.

Практический порядок действий

Очистите градиенты в начале окна, выполните прямой и обратный проходы для каждого нормированного микробатча, затем один шаг оптимизатора. Очистите градиенты для следующего окна. Отдельные прямые проходы обычно не требуют сохранения прежних вычислительных графов. [2] Явно обработайте последнее неполное окно. При смешанной точности сохраняйте постоянный масштаб во время накопления, снимите масштабирование один раз после последнего обратного прохода и затем обрежьте общий градиент перед шагом. Обновляйте scaler на этой границе. [1] Планировщик, заданный на обновление оптимизатора, должен продвигаться только при фактическом обновлении.

Наглядный пример

Допустим, устройство обрабатывает четыре микробатча по восемь примеров. Эффективный батч содержит 32 примера, поэтому каждое среднее получает вес 1/4. Последнее окно с 8, 8, 8 и 2 примерами содержит 26 примеров: веса равны 8/26, 8/26, 8/26 и 2/26. Это арифметический пример, а не измеренный результат обучения.

Ограничения и проверка

Батчевая нормализация использует статистику каждого физического микробатча, поэтому накопление не воспроизводит автоматически один большой батч. [3] Случайные операции и порядок сложения чисел с плавающей точкой тоже могут менять результаты. Отдельная обрезка каждого микробатча обычно меняет общее обновление. Метод сокращает одновременное хранение активаций, но не устраняет веса модели или состояния оптимизатора; дополнительные проходы могут увеличить время работы.

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

Источники

[1] PyTorch: Automatic Mixed Precision examples [2] PyTorch: torch.Tensor.backward [3] PyTorch: BatchNorm1d