Grass: 구조화된 희소 그래디언트를 활용한 계산 효율적 저메모리 LLM 학습

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

June 25, 2024

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

cs.AI

초록

대규모 언어 모델(LLM)의 학습과 미세 조정은 종종 제한된 GPU 메모리로 인해 병목 현상을 겪습니다. 기존의 투영 기반 최적화 방법들은 최적화 상태 메모리를 줄이기 위해 그래디언트를 저차원 부분 공간에 투영함으로써 이 문제를 해결하지만, 일반적으로 밀집된 투영 행렬을 사용하여 계산 및 메모리 오버헤드를 초래할 수 있습니다. 본 연구에서는 희소 투영을 활용하여 그래디언트를 구조화된 희소 업데이트로 변환하는 새로운 접근 방식인 Grass(GRAdient Structured Sparsification)를 제안합니다. 이 설계는 최적화 상태에 대한 메모리 사용량을 크게 줄일 뿐만 아니라 그래디언트 메모리 공간, 계산 및 통신 비용을 최소화하여 상당한 처리량 개선을 이끌어냅니다. 사전 학습 및 미세 조정 작업에 대한 광범위한 실험을 통해 Grass가 전체 랭크 학습 및 기존 투영 기반 방법과 경쟁력 있는 성능을 달성함을 입증했습니다. 특히, Grass는 단일 40GB A100 GPU에서 13B 파라미터 LLaMA 모델의 반정밀도 사전 학습을 가능하게 하며, 이는 이전 방법으로는 불가능했던 성과입니다. 또한 8-GPU 시스템에서 최대 2배의 처리량 개선을 제공합니다. 코드는 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 .