GaLore: Addestramento Efficiente in Memoria di LLM tramite Proiezione a Basso Rango del Gradiente

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

March 6, 2024

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

cs.AI

Abstract

L'addestramento di Large Language Models (LLM) presenta sfide significative in termini di memoria, principalmente a causa delle dimensioni crescenti dei pesi e degli stati dell'ottimizzatore. Approcci comuni per la riduzione della memoria, come l'adattamento a basso rango (LoRA), aggiungono una matrice a basso rango addestrabile ai pesi pre-addestrati congelati in ogni livello, riducendo i parametri addestrabili e gli stati dell'ottimizzatore. Tuttavia, tali approcci generalmente ottengono prestazioni inferiori rispetto all'addestramento con pesi a rango completo sia nelle fasi di pre-addestramento che di fine-tuning, poiché limitano la ricerca dei parametri a un sottospazio a basso rango e alterano le dinamiche di addestramento, e potrebbero inoltre richiedere un avvio a rango completo. In questo lavoro, proponiamo Gradient Low-Rank Projection (GaLore), una strategia di addestramento che consente l'apprendimento a parametri completi ma è più efficiente in termini di memoria rispetto ai comuni metodi di adattamento a basso rango come LoRA. Il nostro approccio riduce l'uso della memoria fino al 65,5% negli stati dell'ottimizzatore, mantenendo sia l'efficienza che le prestazioni per il pre-addestramento sulle architetture LLaMA 1B e 7B con il dataset C4 fino a 19,7 miliardi di token, e per il fine-tuning di RoBERTa sui task GLUE. La nostra versione a 8 bit di GaLore riduce ulteriormente la memoria dell'ottimizzatore fino all'82,5% e la memoria totale di addestramento del 63,3%, rispetto a una baseline BF16. In particolare, dimostriamo, per la prima volta, la fattibilità del pre-addestramento di un modello da 7B su GPU consumer con 24 GB di memoria (ad esempio, NVIDIA RTX 4090) senza l'uso di strategie di parallelismo del modello, checkpointing o offloading.

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.