Grass: Effizientes Training von Low-Memory LLM mit strukturierten spärlichen Gradienten

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

June 25, 2024

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

cs.AI

Zusammenfassung

Das Training und Feintuning großer Sprachmodelle (Large Language Models, LLM) werden oft durch begrenzten GPU-Speicher ausgebremst. Während bestehende auf Projektion basierende Optimierungsmethoden dieses Problem lösen, indem sie Gradienten in einen niedrigdimensionalen Unterraum projizieren, um den Speicherbedarf des Optimierers zu reduzieren, verlassen sie sich in der Regel auf dichte Projektionsmatrizen, die jedoch zusätzlichen Rechenaufwand und Speicherbedarf verursachen können. In dieser Arbeit schlagen wir Grass (GRAdient Strukturierte Sparsifikation) vor, einen neuartigen Ansatz, der auf dünnen Projektionen basiert, um Gradienten in strukturierte, spärliche Aktualisierungen umzuwandeln. Dieses Design reduziert nicht nur signifikant den Speicherverbrauch für Optimiererzustände, sondern minimiert auch den Speicherbedarf für Gradienten, Rechen- und Kommunikationskosten, was zu erheblichen Durchsatzverbesserungen führt. Umfangreiche Experimente zu Vorab- und Feintuning-Aufgaben zeigen, dass Grass eine wettbewerbsfähige Leistung im Vergleich zum Training mit voller Rangordnung und bestehenden auf Projektion basierenden Methoden erzielt. Insbesondere ermöglicht Grass das Vorabtraining eines 13B-Parameter-LLaMA-Modells in Halbpräzision auf einer einzelnen 40GB A100-GPU - eine Leistung, die für frühere Methoden undenkbar war - und führt zu einer bis zu 2-fachen Durchsatzverbesserung auf einem 8-GPU-System. Der Code ist unter https://github.com/aashiqmuhamed/GRASS verfügbar.

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 .