GaLore: 그래디언트 저랭크 투영을 통한 메모리 효율적인 대형 언어 모델 학습

초록

대규모 언어 모델(LLM)을 학습시키는 것은 주로 가중치와 옵티마이저 상태의 크기가 증가함에 따라 상당한 메모리 문제를 야기합니다. 일반적인 메모리 절약 접근법인 저순위 적응(LoRA)은 각 레이어의 고정된 사전 학습 가중치에 학습 가능한 저순위 행렬을 추가하여 학습 가능한 매개변수와 옵티마이저 상태를 줄입니다. 그러나 이러한 접근법은 일반적으로 사전 학습 및 미세 조정 단계에서 전체 순위 가중치를 사용한 학습보다 성능이 떨어지는데, 이는 매개변수 탐색을 저순위 부분공간으로 제한하고 학습 역학을 변경하며, 추가로 전체 순역 웜 스타트가 필요할 수 있기 때문입니다. 본 연구에서는 전체 매개변수 학습을 허용하면서도 LoRA와 같은 일반적인 저순위 적응 방법보다 메모리 효율적인 학습 전략인 Gradient Low-Rank Projection(GaLore)을 제안합니다. 우리의 접근법은 C4 데이터셋에서 최대 19.7B 토큰으로 LLaMA 1B 및 7B 아키텍처를 사전 학습하고, GLUE 작업에서 RoBERTa를 미세 조정하는 동안 옵티마이저 상태 메모리 사용량을 최대 65.5%까지 줄이면서도 효율성과 성능을 유지합니다. 또한, 8비트 GaLore는 BF16 기준선에 비해 옵티마이저 메모리를 최대 82.5%, 전체 학습 메모리를 63.3%까지 추가로 절약합니다. 특히, 우리는 모델 병렬화, 체크포인팅 또는 오프로딩 전략 없이도 24GB 메모리를 가진 소비자용 GPU(예: NVIDIA RTX 4090)에서 7B 모델을 사전 학습하는 것이 가능함을 처음으로 입증했습니다.

English

Training Large Language Models (LLMs) presents significant memory challenges, predominantly due to the growing size of weights and optimizer states. Common memory-reduction approaches, such as low-rank adaptation (LoRA), add a trainable low-rank matrix to the frozen pre-trained weight in each layer, reducing trainable parameters and optimizer states. However, such approaches typically underperform training with full-rank weights in both pre-training and fine-tuning stages since they limit the parameter search to a low-rank subspace and alter the training dynamics, and further, may require full-rank warm start. In this work, we propose Gradient Low-Rank Projection (GaLore), a training strategy that allows full-parameter learning but is more memory-efficient than common low-rank adaptation methods such as LoRA. Our approach reduces memory usage by up to 65.5% in optimizer states while maintaining both efficiency and performance for pre-training on LLaMA 1B and 7B architectures with C4 dataset with up to 19.7B tokens, and on fine-tuning RoBERTa on GLUE tasks. Our 8-bit GaLore further reduces optimizer memory by up to 82.5% and total training memory by 63.3%, compared to a BF16 baseline. Notably, we demonstrate, for the first time, the feasibility of pre-training a 7B model on consumer GPUs with 24GB memory (e.g., NVIDIA RTX 4090) without model parallel, checkpointing, or offloading strategies.

GaLore: 그래디언트 저랭크 투영을 통한 메모리 효율적인 대형 언어 모델 학습

GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection

초록

Support