GaLore: 勾配低ランク射影によるメモリ効率の良いLLMトレーニング

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）があります。これは、各層の凍結された事前学習済み重みに学習可能な低ランク行列を追加し、学習可能なパラメータとオプティマイザ状態を削減します。しかし、このような手法は、パラメータ探索を低ランク部分空間に制限し、トレーニングダイナミクスを変更するため、事前学習とファインチューニングの両方の段階で、フルランク重みを用いたトレーニングに比べて性能が劣ることが一般的です。さらに、フルランクのウォームスタートが必要となる場合もあります。 本研究では、Gradient Low-Rank Projection（GaLore）を提案します。これは、フルパラメータ学習を可能にしつつ、LoRAなどの一般的な低ランク適応手法よりもメモリ効率の高いトレーニング戦略です。我々のアプローチは、LLaMA 1Bおよび7BアーキテクチャでのC4データセットを用いた最大19.7Bトークンの事前学習、およびGLUEタスクでのRoBERTaのファインチューニングにおいて、効率と性能を維持しながら、オプティマイザ状態のメモリ使用量を最大65.5%削減します。さらに、8ビットGaLoreは、BF16ベースラインと比較して、オプティマイザメモリを最大82.5%、トレーニング全体のメモリを63.3%削減します。特に、初めて、24GBメモリを搭載したコンシューマーGPU（例：NVIDIA RTX 4090）上で、モデル並列化、チェックポイント、オフロード戦略なしに7Bモデルの事前学習が可能であることを実証しました。

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.