Градиентный чекпоинтинг: как он экономит память GPU?
Автоматический перевод
Эта статья была автоматически переведена с оригинальной английской версии.
Градиентный чекпоинтинг сохраняет меньше промежуточных активаций при прямом проходе и повторно вычисляет недостающие при обратном. Используйте его, когда из-за активаций обучающий батч или последовательность нужной длины не помещаются в память GPU, а дополнительные вычисления допустимы.
Он сокращает память, занятую сохранёнными активациями. Он не удаляет веса модели, состояние оптимизатора или все временные буферы.
Что вычисляется повторно?
Для вычисления градиентов при обратном проходе нужны промежуточные результаты прямого прохода. Без чекпоинтинга фреймворк сохраняет необходимые результаты. С чекпоинтингом он сохраняет выбранные входные данные или промежуточные результаты и повторно выполняет соответствующие операции, когда нужны их значения.
Например, фреймворк может сохранить вход группы слоёв трансформера и удалить их внутренние активации. При обратном проходе он повторно выполняет эту группу операций, используя сохранённый вход. Сохранённый вход — это чекпоинт активаций. Он отличается от чекпоинта модели, записанного на диск для восстановления.
Исходный анализ Чена и соавторов описывает стратегию с равными сегментами, которая требует O(sqrt(n)) памяти активаций для сети из n слоёв и примерно одного дополнительного прямого прохода на шаг обучения. У реальных трансформеров активации имеют разные размеры, а реализации аттеншна и правила фреймворков различаются. Измеряйте фактические затраты памяти и времени, а не используйте асимптотическую оценку как прогноз для деплоя.
Включите интеграцию, предусмотренную вашим стеком
Для моделей, поддерживаемых Hugging Face Trainer, настройка gradient_checkpointing=True включает общий механизм. При использовании FSDP проверьте вместо этого интеграцию чекпоинтинга активаций во фреймворке. Текущая документация Transformers Trainer рекомендует activation_checkpointing для FSDP, поскольку обычный градиентный чекпоинтинг может вызывать дополнительную операцию all-gather при обратном проходе.
Сравните одну и ту же конфигурацию обучения с чекпоинтингом и без него. Запишите пиковый объём выделенной памяти GPU, число примеров или токенов в секунду и проверьте, помещается ли последовательность нужной длины. Также проверьте, остаются ли градиенты корректными, если модель использует пользовательские операции или обучение адаптеров.
Если большую часть памяти занимает состояние оптимизатора, рассмотрите шардирование состояния или настройку с малым числом обучаемых параметров; повторное вычисление активаций не освободит эту память.
Раздел о градиентном чекпоинтинге в LLM Engineering Guide объясняет расчёт по сегментам.