GaLore: обучение больших моделей на потребительских GPU

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

GaLore — это новый метод оптимизации, представленный исследователями из Университета Карнеги-Меллона и Google, который позволяет обучать большие языковые модели на обычных потребительских GPU. Вместо того чтобы использовать полные матрицы градиентов, GaLore проецирует их в низкоранговое пространство, что значительно снижает требования к памяти. В тестах модель с 7 миллиардами параметров обучалась на GPU с 24 ГБ памяти — ранее для этого требовалось не менее 40 ГБ.

Этот подход решает ключевую проблему, с которой сталкиваются исследователи и разработчики: стоимость и доступность оборудования для обучения современных моделей. Потребительские видеокарты, такие как NVIDIA RTX 4090, становятся способными обучать модели, которые раньше были прерогативой крупных дата-центров. Это может демократизировать доступ к разработке ИИ, позволяя большему числу людей участвовать в создании передовых технологий.

Предыстория и контекст

Традиционные методы оптимизации, такие как Adam, требуют хранения полных градиентов и моментов, что приводит к огромному потреблению памяти при обучении больших моделей. Например, для модели с 7 миллиардами параметров требуется около 112 ГБ памяти только для градиентов и моментов, что делает обучение на обычном оборудовании невозможным. Техники, такие как смешанная точность и распределенное обучение, помогают, но имеют свои ограничения.

GaLore основан на идее, что градиенты в процессе обучения часто имеют низкий ранг, то есть их можно эффективно представить в пространстве меньшей размерности. Используя этот факт, метод проецирует градиенты в низкоранговое пространство с помощью случайной проекции, обновляет параметры там, а затем проецирует обратно. Это снижает потребление памяти с O(N) до O(N·r), где r — ранг проекции, который обычно значительно меньше N.

Интересно, что GaLore не требует предварительного вычисления SVD (сингулярного разложения) и может работать с фиксированным рангом, что упрощает его использование. В отличие от методов, таких как LoRA, которые модифицируют архитектуру модели, GaLore работает на уровне оптимизатора, что делает его универсальным и легким в интеграции с существующими моделями.

Как GaLore работает?

GaLore использует случайную проекцию для снижения размерности градиентов. На каждом шаге оптимизации градиенты проецируются в низкоранговое пространство, где применяется стандартный оптимизатор (например, Adam), а затем обновления проецируются обратно в исходное пространство параметров. Это позволяет значительно сократить объем памяти, необходимый для хранения градиентов и моментов, не теряя при этом качества обучения.

Технические подробности

В экспериментах исследователи обучали модель с 7 миллиардами параметров на GPU NVIDIA RTX 4090 с 24 ГБ памяти. Для сравнения, стандартный Adam требовал бы более 112 ГБ памяти, что делает обучение невозможным на таком оборудовании. С GaLore потребление памяти снизилось до 24 ГБ, что позволило успешно завершить обучение без потери качества.

Они также протестировали GaLore на моделях меньшего размера, таких как 1B и 2B параметров, и обнаружили, что метод работает стабильно и сходится быстрее, чем стандартные подходы, при сохранении точности. В частности, модель LLaMA-7B, обученная с помощью GaLore, показала сравнимые результаты с моделью, обученной на более мощном оборудовании.

Важно отметить, что GaLore не требует изменений в архитектуре модели или специального кода для конкретных задач. Он интегрируется с популярными библиотеками, такими как PyTorch, и может быть использован с любым оптимизатором, поддерживающим низкоранговое обновление. Это делает его доступным для широкого круга разработчиков.

Кого затронет и как

Для исследователей и разработчиков, работающих в области ИИ, GaLore открывает новые горизонты: теперь можно обучать большие модели без необходимости арендовать дорогие облачные кластеры. Это особенно важно для стартапов и академических групп с ограниченным бюджетом. Кроме того, метод может быть использован для тонкой настройки больших моделей на потребительском оборудовании, что упрощает персонализацию моделей для конкретных задач.

Для сообщества с открытым исходным кодом GaLore означает, что больше людей смогут участвовать в разработке и обучении передовых моделей, что может ускорить инновации и сделать ИИ более доступным. В России и СНГ это также актуально, учитывая сложности с импортом высокопроизводительного оборудования и ограничения на облачные сервисы.

Что будет дальше

Исследовательская группа планирует продолжить работу над улучшением GaLore, включая оптимизацию ранга проекции и адаптацию метода для других типов моделей, таких как трансформеры для зрения. Также ожидается, что GaLore будет интегрирован в популярные фреймворки, такие как Hugging Face Transformers, что упростит его использование.

В ближайшие месяцы мы можем ожидать появления новых статей и примеров, демонстрирующих применение GaLore в различных сценариях. Это может привести к появлению новых инструментов и библиотек, которые сделают обучение больших моделей еще более доступным.

Итог

GaLore — это значительный шаг вперед в области оптимизации обучения больших моделей. Он позволяет обучать модели с миллиардами параметров на обычных потребительских GPU, что открывает новые возможности для исследователей и разработчиков по всему миру. Следите за развитием этой технологии — она может изменить ландшафт искусственного интеллекта, сделав его более демократичным и доступным.