GaLore: Speicher-effizientes Training von LLM durch Gradienten-Niedrig-Rang-Projektion.

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

March 6, 2024

papers.authors: Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, Yuandong Tian

cs.AI

papers.abstract

Das Training von großen Sprachmodellen (LLMs) birgt erhebliche Speicherherausforderungen, hauptsächlich aufgrund der wachsenden Größe der Gewichte und Optimizer-Zustände. Gängige Speicherreduktionsansätze wie die Low-Rank-Anpassung (LoRA) fügen jeder Schicht eine trainierbare Low-Rank-Matrix zu den eingefrorenen vortrainierten Gewichten hinzu, um die trainierbaren Parameter und Optimizer-Zustände zu reduzieren. Solche Ansätze erzielen jedoch in der Regel eine geringere Leistung als das Training mit vollständigen Gewichten sowohl in der Vortrainings- als auch in den Feinabstimmungsphasen, da sie die Parametersuche auf einen Low-Rank-Unterraum beschränken und die Trainingsdynamik verändern und möglicherweise einen Warmstart mit vollständigen Gewichten erfordern. In dieser Arbeit schlagen wir Gradienten-Low-Rank-Projektion (GaLore) vor, eine Trainingsstrategie, die ein vollständiges Lernen der Parameter ermöglicht, aber speichereffizienter ist als gängige Low-Rank-Anpassungsmethoden wie LoRA. Unser Ansatz reduziert den Speicherverbrauch der Optimizer-Zustände um bis zu 65,5%, während Effizienz und Leistungsfähigkeit sowohl beim Vortraining auf LLaMA 1B und 7B Architekturen mit dem C4-Datensatz mit bis zu 19,7B Tokens als auch bei der Feinabstimmung von RoBERTa auf GLUE-Aufgaben erhalten bleiben. Unser 8-Bit GaLore reduziert den Speicherbedarf des Optimizers um bis zu 82,5% und den Gesamtspeicherbedarf des Trainings um 63,3% im Vergleich zu einem BF16-Baseline. Bemerkenswert ist, dass wir erstmals die Machbarkeit des Vortrainings eines 7B-Modells auf Consumer-GPUs mit 24GB Speicher (z. B. NVIDIA RTX 4090) ohne Modellparallelität, Checkpointing oder Auslagerungsstrategien demonstrieren.

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.