GaLore: Обучение памяти-эффективной модели языкового моделирования с помощью проекции градиента низкого ранга

GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection

March 6, 2024

Авторы: Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, Yuandong Tian

cs.AI

Аннотация

Обучение больших языковых моделей (LLM) представляет существенные проблемы с памятью, в основном из-за растущего размера весов и состояний оптимизатора. Общие подходы к сокращению памяти, такие как адаптация низкого ранга (LoRA), добавляют обучаемую матрицу низкого ранга к замороженным предварительно обученным весам в каждом слое, сокращая обучаемые параметры и состояния оптимизатора. Однако такие подходы обычно производят результаты ниже, чем обучение с весами полного ранга как на этапе предварительного обучения, так и на этапе донастройки, поскольку они ограничивают поиск параметров до низкорангового подпространства и изменяют динамику обучения, и, кроме того, могут потребовать начальной настройки с полным рангом. В данной работе мы предлагаем Градиентную Проекцию Низкого Ранга (GaLore), стратегию обучения, которая позволяет полное обучение параметров, но более эффективна с точки зрения использования памяти, чем общие методы адаптации низкого ранга, такие как LoRA. Наш подход снижает использование памяти на оптимизаторы до 65.5%, сохраняя при этом эффективность и производительность как для предварительного обучения на архитектурах LLaMA 1B и 7B с набором данных C4 объемом до 19.7B токенов, так и для донастройки RoBERTa на задачах GLUE. Наш 8-битный GaLore дополнительно снижает использование памяти оптимизатора до 82.5% и общую память обучения на 63.3%, по сравнению с базовым уровнем BF16. Следует отметить, что мы демонстрируем впервые возможность предварительного обучения модели 7B на потребительских GPU с 24 ГБ памяти (например, NVIDIA RTX 4090) без параллельного моделирования, стратегий контрольных точек или выгрузки.

English

Training Large Language Models (LLMs) presents significant memory challenges, predominantly due to the growing size of weights and optimizer states. Common memory-reduction approaches, such as low-rank adaptation (LoRA), add a trainable low-rank matrix to the frozen pre-trained weight in each layer, reducing trainable parameters and optimizer states. However, such approaches typically underperform training with full-rank weights in both pre-training and fine-tuning stages since they limit the parameter search to a low-rank subspace and alter the training dynamics, and further, may require full-rank warm start. In this work, we propose Gradient Low-Rank Projection (GaLore), a training strategy that allows full-parameter learning but is more memory-efficient than common low-rank adaptation methods such as LoRA. Our approach reduces memory usage by up to 65.5% in optimizer states while maintaining both efficiency and performance for pre-training on LLaMA 1B and 7B architectures with C4 dataset with up to 19.7B tokens, and on fine-tuning RoBERTa on GLUE tasks. Our 8-bit GaLore further reduces optimizer memory by up to 82.5% and total training memory by 63.3%, compared to a BF16 baseline. Notably, we demonstrate, for the first time, the feasibility of pre-training a 7B model on consumer GPUs with 24GB memory (e.g., NVIDIA RTX 4090) without model parallel, checkpointing, or offloading strategies.