Grass: Rekenkundig efficiënte LLM-training met laag geheugengebruik via gestructureerd sparse gradients

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

June 25, 2024

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

cs.AI

Samenvatting

De training en fine-tuning van grote taalmodellen (LLM's) worden vaak beperkt door het beperkte GPU-geheugen. Hoewel bestaande optimalisatiemethoden op basis van projecties dit aanpakken door gradiënten te projecteren in een lager-dimensionale deelruimte om het geheugengebruik van de optimalisatietoestand te verminderen, vertrouwen ze meestal op dichte projectiematrices, wat rekenkundige en geheugenoverhead kan introduceren. In dit werk stellen we Grass (GRAdient Structured Sparsification) voor, een nieuwe aanpak die gebruikmaakt van sparse projecties om gradiënten om te zetten in gestructureerde sparse updates. Dit ontwerp vermindert niet alleen aanzienlijk het geheugengebruik voor optimalisatietoestanden, maar minimaliseert ook de geheugenvoetafdruk van gradiënten, rekenkosten en communicatiekosten, wat leidt tot aanzienlijke doorvoerverbeteringen. Uitgebreide experimenten op taken voor pretraining en fine-tuning tonen aan dat Grass competitieve prestaties levert in vergelijking met training op volledige rang en bestaande projectiegebaseerde methoden. Opmerkelijk is dat Grass halfprecisie-pretraining van een LLaMA-model met 13B parameters mogelijk maakt op een enkele 40GB A100 GPU—een prestatie die voor eerdere methoden onhaalbaar was—en een tot 2 keer hogere doorvoer oplevert op een 8-GPU-systeem. Code is te vinden op 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 .