GaLore: Geheugenefficiënte LLM-training door Gradient Low-Rank Projectie

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

March 6, 2024

Auteurs: Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, Yuandong Tian

cs.AI

Samenvatting

Het trainen van Large Language Models (LLMs) brengt aanzienlijke geheugenuitdagingen met zich mee, voornamelijk door de toenemende omvang van gewichten en optimizer-statussen. Gangbare benaderingen om geheugen te reduceren, zoals low-rank adaptation (LoRA), voegen een trainbare low-rank matrix toe aan de bevroren voorgetrainde gewichten in elke laag, waardoor het aantal trainbare parameters en optimizer-statussen wordt verminderd. Dergelijke benaderingen presteren echter doorgaans minder goed dan training met full-rank gewichten, zowel in de voorafgaande training als in de fine-tuning-fase, omdat ze de parameterzoekruimte beperken tot een low-rank deelruimte en de trainingsdynamiek veranderen, en bovendien mogelijk een full-rank warm start vereisen. In dit werk stellen we Gradient Low-Rank Projection (GaLore) voor, een trainingsstrategie die volledige parameterleren mogelijk maakt, maar toch geheugenefficiënter is dan gangbare low-rank adaptatiemethoden zoals LoRA. Onze aanpak reduceert het geheugengebruik van optimizer-statussen met tot wel 65,5% terwijl zowel efficiëntie als prestaties behouden blijven voor voorafgaande training op LLaMA 1B en 7B architecturen met de C4 dataset met tot 19,7B tokens, en voor fine-tuning van RoBERTa op GLUE-taken. Onze 8-bit GaLore reduceert het geheugen van de optimizer verder met tot 82,5% en het totale trainingsgeheugen met 63,3% in vergelijking met een BF16-basislijn. Opmerkelijk is dat we voor het eerst de haalbaarheid aantonen van het vooraf trainen van een 7B-model op consumenten-GPU's met 24GB geheugen (bijv. NVIDIA RTX 4090) zonder modelparallelisme, checkpointing of offloadingstrategieën.

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.