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

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는 LLaMA 7B 모델을 24GB GPU에서 훈련할 수 있도록 하며, 최적화 상태 메모리를 최대 65.5% 절감하는 기울기 저랭크 투영 기반 훈련 전략이다.

핵심 기여도

GaLore는 LoRA 대비 최적화 상태 메모리를 최대 65.5% 절감하면서 LLaMA 1B/7B 모델의 사전 훈련과 GLUE 태스크에서 RoBERTa 미세조정 성능을 유지.
8-bit GaLore는 BF16 기준으로 최적화 메모리를 82.5%, 전체 훈련 메모리를 63.3% 절감.
GaLore는 ReLoRA와 달리 완전 랭크 워밍업 없이 7B 모델을 24GB GPU에서 훈련 가능.
GLUE 태스크에서 RoBERTa-Base에 적용한 GaLore는 LoRA 대비 85.61 대 85.89의 평균 점수를 달성.

핵심 아이디어

기존 LoRA는 가중치 행렬을 저랭크로 재매개변수화하여 메모리를 절감하지만, 학습 동역학을 제한하고 성능 저하를 유발할 수 있다. GaLore는 이와 달리, 가중치 행렬이 아닌 **기울기 행렬** $ G \in \mathbb{R}^{m \times n} $의 저랭크 구조를 활용한다. 이는 기울기의 변화가 느리고 저랭크로 유지된다는 이론적 근거를 바탕으로 한다. GaLore는 기울기 행렬 $ G $를 $ P^\top G Q $ 형태로 투영하는 두 개의 저랭크 행렬 $ P \in \mathbb{R}^{m \times r} $, $ Q \in \mathbb{R}^{n \times r} $를 사용하며, 이 투영 행렬은 200 iteration마다 간헐적으로 업데이트된다. 이 방식은 최적화 상태 메모리를 줄이면서도 전체 파라미터 학습을 유지할 수 있다.

기술적 접근법

**GaLore 알고리즘**: 기울기 행렬 $ G $를 $ P^\top G Q $로 투영. $ P, Q $는 $ r $-랭크 행렬로, $ r \ll \min(m, n) $.
**메모리 절감**: 최적화 상태 메모리가 $ 3mn $에서 $ 3r(m+n) $로 줄어들며, $ r $이 작을수록 절감 효과 증대.
**8-bit GaLore**: 8-bit 최적화기와 레이어별 가중치 업데이트 기법을 결합하여 메모리 절감.
**적용 가능한 최적화기**: AdamW, 8-bit Adam, Adafactor 등과 호환 가능.
**하이퍼파라미터**: 투영 행렬 $ P, Q $는 200 iteration마다 갱신되며, 학습률은 각 방법에 맞게 조정.

주요 결과

**LLaMA 7B 사전 훈련**: C4 데이터셋(19.7B 토큰)에서 GaLore는 최적화 상태 메모리를 65.5% 절감하면서 성능 유지.
**RoBERTa 미세조정**: GLUE 태스크에서 LoRA(85.61) 대비 GaLore(85.89)가 0.28% 높은 평균 점수 달성.
**8-bit GaLore**: BF16 대비 최적화 메모리 82.5%, 전체 훈련 메모리 63.3% 절감.
**소비자 GPU 훈련**: 24GB 메모리의 NVIDIA RTX 4090에서 LLaMA 7B 모델을 모델 병렬, 체크포인팅, 오프로딩 없이 훈련 가능.

의의 및 한계

GaLore는 기존 저랭크 조정법(LoRA)의 성능 제약을 극복하면서도 메모리 효율성을 유지하는 새로운 접근법이다. 기울기 투영을 통해 전체 파라미터 학습이 가능하며, 기존 최적화기와 호환되어 쉽게 적용할 수 있다. 특히, 24GB GPU에서 7B 규모 모델을 훈련하는 것은 기존 기술에서는 불가능했던 점에서 혁신적인 의의를 가진다. 다만, 저랭크 투영 행렬 $ P, Q $의 갱신 주기나 랭크 $ r $의 선택은 성능에 영향을 줄 수 있으며, 이에 대한 최적화는 추가 연구가 필요하다. 또한, 비전 트랜스포머나 확산 모델 등 다른 모델 구조에의 적용 가능성도 개방 문제로 남아 있다.

실용적 활용

GaLore는 메모리 제약이 있는 소비자 GPU 환경에서 대형 언어 모델 훈련을 가능하게 하며, 연구실이나 중소 기업에서도 저비용으로 LLM 개발이 가능하도록 지원한다. 또한, GLUE와 같은 자연어 이해 태스크에서 기존 저랭크 조정법 대비 더 높은 성능을 보이는 만큼, 미세조정 기반의 실제 응용(예: 챗봇, 문서 분류)에도 유용하게 활용될 수 있다.