APOLLO: Эффективное обучение больших языковых моделей с использованием памяти

В последние годы большие языковые модели (LLMs) продемонстрировали выдающиеся способности в различных областях, таких как обработка естественного языка, создание текста и даже в задачах, связанных с пониманием контекста. Однако обучение этих моделей требует значительных вычислительных ресурсов и памяти, что создает барьеры для многих исследователей и разработчиков. В данной статье мы рассмотрим новую методику, предложенную в работе "APOLLO: SGD-like Memory, AdamW-level Performance", которая направлена на оптимизацию использования памяти при обучении LLM, сохраняя при этом высокую производительность.

Большие языковые модели, такие как LLaMA и GPT, достигли впечатляющих результатов благодаря увеличению количества параметров до миллиардов. Однако их обучение требует огромных вычислительных ресурсов, что делает его доступным только для крупных организаций. Например, обучение модели LLaMA-7B требует более 58 ГБ оперативной памяти, из которых 28 ГБ используется для состояний оптимизатора AdamW. Для моделей с еще большим количеством параметров, таких как GPT-3, требования к памяти могут достигать 700 ГБ, что создает колоссальные трудности при обучении.

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

Проблема и ее решение

Высокая память при использовании AdamW

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

Подход APOLLO

APOLLO предлагает новое решение, которое заключается в том, чтобы уменьшить избыточность в правилах адаптации скорости обучения AdamW. Мы обнаружили, что скорость обучения может быть структурирована в виде обновлений на уровне каналов или тензоров, что позволяет значительно снизить использование памяти. В результате мы разработали метод, который использует вспомогательное состояние оптимизатора с низким рангом, основанное на случайных проекциях, для приближения структурированных факторов скорости обучения.

Структурированные обновления скорости обучения

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

Эксперименты и результаты

Эффективность APOLLO

Мы провели обширные эксперименты с различными архитектурами моделей и задачами, чтобы продемонстрировать эффективность APOLLO. Результаты показали, что APOLLO и его мини-версия (APOLLO-Mini) могут достигать производительности на уровне AdamW, при этом значительно снижая использование памяти. Например, APOLLO позволяет обучать модель LLaMA-7B на одном GPU с менее чем 12 ГБ памяти, что делает обучение доступным даже для исследователей с ограниченными ресурсами.

Сравнение с другими методами

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

Улучшение производительности

Эксперименты показали, что APOLLO и APOLLO-Mini обеспечивают улучшение производительности в задачах предобучения и дообучения. В частности, APOLLO достигает до 2.8 раз снижения перплексии в валидации по сравнению с AdamW, сохраняя при этом значительно меньшие затраты памяти.

Заключение

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