Grass: Computação Eficiente de Treinamento LLM com Baixa Memória usando Gradientes Esparsos Estruturados

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

Resumo

O treinamento e ajuste fino de grandes modelos de linguagem (LLM) frequentemente são limitados pela memória GPU disponível. Enquanto os métodos de otimização baseados em projeção existentes lidam com isso projetando gradientes em um subespaço de menor dimensão para reduzir a memória do estado do otimizador, eles geralmente dependem de matrizes de projeção densas, o que pode introduzir sobrecargas computacionais e de memória. Neste trabalho, propomos Grass (GRAdient Stuctured Sparsification), uma abordagem inovadora que utiliza projeções esparsas para transformar gradientes em atualizações esparsas estruturadas. Esse design não apenas reduz significativamente o uso de memória para os estados do otimizador, mas também minimiza a pegada de memória do gradiente, os custos de computação e comunicação, resultando em melhorias substanciais de rendimento. Experimentos extensivos em tarefas de pré-treinamento e ajuste fino demonstram que o Grass alcança desempenho competitivo em relação ao treinamento de classificação total e aos métodos existentes baseados em projeção. Notavelmente, o Grass possibilita o pré-treinamento em meia precisão de um modelo LLaMA de 13B parâmetros em uma única GPU A100 de 40GB - uma conquista inviável para métodos anteriores - e resulta em uma melhoria de rendimento de até 2 vezes em um sistema com 8 GPUs. O código pode ser encontrado em 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 .