Grass: Addestramento Efficiente di LLM a Basso Consumo di Memoria con Gradienti Strutturalmente Sparsi

Grass: Compute Efficient Low-Memory LLM Training with Structured Sparse Gradients

June 25, 2024

Autori: Aashiq Muhamed, Oscar Li, David Woodruff, Mona Diab, Virginia Smith

cs.AI

Abstract

L'addestramento e il fine-tuning di grandi modelli linguistici (LLM) sono spesso limitati dalla memoria GPU disponibile. Sebbene i metodi di ottimizzazione basati su proiezione esistenti affrontino questo problema proiettando i gradienti in un sottospazio a dimensione ridotta per diminuire la memoria degli stati dell'ottimizzatore, essi si basano tipicamente su matrici di proiezione dense, che possono introdurre sovraccarichi computazionali e di memoria. In questo lavoro, proponiamo Grass (GRAdient Structured Sparsification), un approccio innovativo che sfrutta proiezioni sparse per trasformare i gradienti in aggiornamenti strutturati e sparsi. Questo design non solo riduce significativamente l'uso della memoria per gli stati dell'ottimizzatore, ma minimizza anche l'impronta di memoria dei gradienti, i costi computazionali e di comunicazione, portando a sostanziali miglioramenti della produttività. Esperimenti estesi su compiti di pre-addestramento e fine-tuning dimostrano che Grass raggiunge prestazioni competitive rispetto all'addestramento a rango pieno e ai metodi basati su proiezione esistenti. In particolare, Grass consente il pre-addestramento in mezza precisione di un modello LLaMA da 13B parametri su una singola GPU A100 da 40GB—un risultato impossibile per i metodi precedenti—e produce un miglioramento fino a 2 volte della produttività su un sistema a 8 GPU. Il codice è disponibile all'indirizzo https://github.com/aashiqmuhamed/GRASS.

English

Large language model (LLM) training and finetuning are often bottlenecked by limited GPU memory. While existing projection-based optimization methods address this by projecting gradients into a lower-dimensional subspace to reduce optimizer state memory, they typically rely on dense projection matrices, which can introduce computational and memory overheads. In this work, we propose Grass (GRAdient Stuctured Sparsification), a novel approach that leverages sparse projections to transform gradients into structured sparse updates. This design not only significantly reduces memory usage for optimizer states but also minimizes gradient memory footprint, computation, and communication costs, leading to substantial throughput improvements. Extensive experiments on pretraining and finetuning tasks demonstrate that Grass achieves competitive performance to full-rank training and existing projection-based methods. Notably, Grass enables half-precision pretraining of a 13B parameter LLaMA model on a single 40GB A100 GPU--a feat infeasible for previous methods--and yields up to a 2times throughput improvement on an 8-GPU system. Code can be found at https://github.com/aashiqmuhamed/GRASS .