Обучение со смешанной точностью и AMP

При обучении со смешанной точностью большинство операций выполняется с 16-битными числами с плавающей точкой (FP16), а не с 32-битными (FP32). FP16 использует вдвое меньше памяти и работает быстрее на тензорных ядрах современных GPU, тогда как несколько чувствительных к точности операций для стабильности остаются в FP32.

Смешанная точность обеспечивает:

• экономию памяти примерно в 2 раза, что позволяет использовать более крупные пакеты или модели
• Более быстрое умножение матриц на тензорных ядрах
• Незначительное или нулевое снижение итоговой точности модели при корректном применении

FP16 имеет узкий диапазон представимых значений. Небольшие значения градиентов могут обнулиться из-за выхода за нижнюю границу диапазона, незаметно остановив обучение. Это основная проблема, которую должна решать смешанная точность, и причина, по которой нельзя просто преобразовать всё в FP16.

Автоматическая смешанная точность (AMP) в PyTorch обрабатывает эти детали с помощью двух инструментов: autocast выбирает точность для каждой операции, а GradScaler предотвращает потерю значимости градиентов.

Оберните прямой проход в autocast(). Внутри него PyTorch автоматически выполняет каждую операцию с наиболее безопасной точностью: умножение матриц — в FP16, а операции сокращения, такие как softmax и вычисление функции потерь, — в FP32.

GradScaler борется с потерей значимости градиентов, умножая функцию потерь на большой коэффициент масштабирования перед обратным распространением ошибки. Это переносит малые градиенты в диапазон, представимый в FP16; перед шагом оптимизатора масштабирование отменяется.

scaler.scale(loss).backward() увеличивает масштаб функции потерь, а затем запускает обратное распространение ошибки. Получившиеся градиенты также масштабируются, что не позволяет малым значениям исчезнуть.

scaler.step(optimizer) сначала возвращает градиентам их истинный масштаб, а затем вызывает optimizer.step(), но только если не появились бесконечности или NaNs. Если произошло переполнение градиентов, шаг пропускается.

scaler.update() корректирует коэффициент масштабирования для следующей итерации: увеличивает его при успешном выполнении шагов и уменьшает после переполнения. Такое адаптивное масштабирование автоматически поддерживает стабильность обучения.

Объединив все элементы, мы получаем полный шаг обучения со смешанной точностью.

Помните о следующем:

• Вызывайте backward() только для масштабированной функции потерь
• autocast оборачивает только прямой проход, но не обратное распространение ошибки и не шаг оптимизатора
• AMP особенно эффективен на GPU с тензорными ядрами; на старом оборудовании прирост меньше

Проверьте свои знания об AMP.

Вы изучили обучение со смешанной точностью с помощью AMP:

• autocast() выбирает FP16 или FP32 для каждой операции прямого прохода
• GradScaler масштабирует функцию потерь, чтобы предотвратить потерю значимости градиентов
• Цикл состоит из scaler.scale(loss).backward(), scaler.step(optimizer), scaler.update()
• Результат: экономия памяти примерно в 2 раза и более быстрое обучение