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

TL;DR

Grass introduces a structured sparse gradient projection method that reduces memory and computational costs in large language model training, enabling efficient pretraining of large models on limited hardware with significant throughput gains.

Contribution

The paper presents Grass, a novel sparse projection technique that improves memory efficiency and training throughput for large language models, outperforming dense and existing projection methods.

Findings

Enables pretraining of 13B LLaMA on a single 40GB GPU

Achieves up to 2x throughput improvement on 8-GPU systems

Maintains competitive performance with full-rank training

Abstract

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 $2\times$ throughput improvement on an 8-GPU system. Code can be found at https://github.com/aashiqmuhamed/GRASS .