Grass: Entrenamiento eficiente de LLM con bajo consumo de memoria mediante gradientes dispersos estructurados

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

June 25, 2024

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

cs.AI

Resumen

El entrenamiento y ajuste fino de modelos de lenguaje de gran escala (LLM, por sus siglas en inglés) suelen verse limitados por la memoria limitada de las GPU. Si bien los métodos de optimización basados en proyección existentes abordan este problema al proyectar los gradientes en un subespacio de menor dimensión para reducir la memoria del estado del optimizador, generalmente dependen de matrices de proyección densas, lo que puede introducir sobrecargas computacionales y de memoria. En este trabajo, proponemos Grass (GRAdient Structured Sparsification), un enfoque novedoso que aprovecha proyecciones dispersas para transformar los gradientes en actualizaciones dispersas estructuradas. Este diseño no solo reduce significativamente el uso de memoria para los estados del optimizador, sino que también minimiza la huella de memoria de los gradientes, los costos de cálculo y de comunicación, lo que conduce a mejoras sustanciales en el rendimiento. Experimentos exhaustivos en tareas de preentrenamiento y ajuste fino demuestran que Grass logra un rendimiento competitivo en comparación con el entrenamiento de rango completo y los métodos basados en proyección existentes. Es notable que Grass permite el preentrenamiento en precisión media de un modelo LLaMA de 13B parámetros en una sola GPU A100 de 40 GB, un logro inviable para métodos anteriores, y ofrece una mejora de hasta 2 veces en el rendimiento en un sistema de 8 GPU. El código está disponible en 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 .