Grass: Compute Efficient Low-Memory LLM Training with Structured Sparse
  Gradients
Authors: name: Aashiq Muhamed; name: Oscar Li; name: David Woodruff; name: Mona Diab; name: Virginia Smith
Published: 2024-06-25 15:50:32

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