GaLore가 그래디언트의 저랭크 구조를 투영하여 옵티마이저 상태 메모리를 82.5% 감소시키고 소비자용 GPU에서 70억 파라미터 모델 학습 가능

Context

대규모 언어 모델 학습 시 옵티마이저 상태(특히 Adam)가 전체 메모리 풋프린트의 상당 부분을 차지하여 소비자용 GPU에서는 수십억 규모의 파라미터를 가진 모델 학습이 불가능했다.

Technical Solution

• 그래디언트를 저랭크 부분공간으로 투영: 깊은 신경망의 그래디언트 내재 저랭크 구조를 이용하여 저장 및 처리 데이터의 차원성 감소
• 동적 부분공간 전환 메커니즘 도입: 학습 과정에서 그래디언트의 저랭크 구조 변화에 응답하여 다양한 저랭크 부분공간 간 동적 전환으로 전체 파라미터 학습 용량 보존
• 8비트 정밀도 옵티마이저와 결합: 그래디언트, 가중치, 옵티마이저 상태를 32비트 부동소수점에서 8비트 정수로 양자화하여 추가 메모리 절감
• 레이어별 가중치 업데이트 구현: PyTorch post-accumulation 훅으로 모든 레이어의 그래디언트를 메모리에 유지하는 대신 한 번에 한 레이어씩 가중치 업데이트하여 메모리 풋프린트 추가 감소
• bitsandbytes 라이브러리에 통합: 양자화 및 전문화된 행렬 매개변수화와의 긴밀한 통합으로 추가 메모리 감소 가능성 탐색

Impact

옵티마이저 상태 저장 시 메모리 82.5% 이상 감소. NVIDIA RTX 4090 같은 소비자용 GPU에서 Llama 기반 70억 파라미터 모델 학습 가능.

Key Takeaway

그래디언트의 저랭크 특성을 활용한 투영 기법과 양자화를 결합하면 고사양 하드웨어 없이도 대규모 모델 학습을 민주화할 수 있으며, 부분공간 전환으로 저차원화의 표현력 한계를 극복할 수 있다.