Grass : Entraînement efficace en calcul et faible mémoire pour les LLM avec des gradients structurés et parcimonieux

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

Résumé

L'entraînement et le réglage fin des grands modèles de langage (LLM) sont souvent limités par la mémoire GPU disponible. Bien que les méthodes d'optimisation basées sur la projection existantes abordent ce problème en projetant les gradients dans un sous-espace de dimension inférieure pour réduire la mémoire des états de l'optimiseur, elles reposent généralement sur des matrices de projection denses, ce qui peut introduire des surcharges de calcul et de mémoire. Dans ce travail, nous proposons Grass (GRAdient Structured Sparsification), une approche novatrice qui exploite des projections creuses pour transformer les gradients en mises à jour structurées et creuses. Cette conception réduit non seulement de manière significative l'utilisation de la mémoire pour les états de l'optimiseur, mais minimise également l'empreinte mémoire des gradients, les coûts de calcul et de communication, conduisant à des améliorations substantielles du débit. Des expériences approfondies sur des tâches de pré-entraînement et de réglage fin démontrent que Grass atteint des performances comparables à l'entraînement en rang complet et aux méthodes de projection existantes. Notamment, Grass permet le pré-entraînement en demi-précision d'un modèle LLaMA de 13 milliards de paramètres sur un seul GPU A100 de 40 Go—une prouesse impossible pour les méthodes précédentes—et offre jusqu'à une amélioration du débit par un facteur de 2 sur un système à 8 GPU. Le code est disponible à l'adresse suivante : 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 .