Gefen: новый оптимизатор для глубокого обучения снижает потребление памяти в 8 раз
Исследователи опубликовали на arXiv препринт, в котором представили новый оптимизатор Gefen для глубокого обучения. Разработка нацелена на снижение потребления памяти, которое является узким местом при тренировке крупных моделей.
Популярный оптимизатор AdamW хранит состояния первого и второго моментов для каждого параметра, что добавляет два буфера размером с параметры. Для модели с миллиардом параметров это около 8 ГБ дополнительной памяти. Gefen решает эту проблему, автоматически разделяя оценки второго момента между блоками параметров и квантуя первый момент с помощью обучаемого кодового словаря.
По словам авторов, Gefen сокращает потребление памяти AdamW примерно в 8 раз, что соответствует уменьшению на 6,5 ГБ на миллиард параметров, при этом производительность остается на уровне AdamW.
Метод мотивирован теоретическим результатом: большие смешанные элементы гессиана ограничивают отношение квадратов градиентов к единице. Это означает, что параметры, выровненные по гессиану, являются естественными кандидатами для совместного использования статистики второго момента. Поскольку вычисление гессиана в масштабе непрактично, Gefen определяет структуру блоков по начальным квадратам градиентов, не требуя метаданных об архитектуре или дополнительных гиперпараметров.
Для квантования первого момента используется точный гистограммный алгоритм динамического программирования, а те же блоки применяются для масштабирования. В разнообразных экспериментах Gefen достиг наименьшего пикового потребления памяти среди сравниваемых методов, сохранив производительность AdamW.
При тренировке с FSDP и DDP снижение памяти позволяет использовать более крупные микробатчи и значительно увеличивает пропускную способность по сравнению с AdamW. Разработчики предоставили полную реализацию на Python, включая CUDA-ядра, на GitHub.
Gefen может стать практичной заменой AdamW с меньшим потреблением памяти, что позволяет тренировать более крупные модели или увеличивать размер батча без дополнительных затрат.
