GaLore: 소비자 등급 하드웨어에서 대규모 모델 훈련 진전

GaLore는 NVIDIA RTX 4090과 같은 소비자 등급 하드웨어에서 최대 70억 파라미터를 가진 대형 언어 모델(LLM) 훈련을 가능하게 하며, 옵티마이저 상태의 메모리 사용량을 크게 줄입니다. 이러한 AI 연구의 민주화는 실무자들이 고성능 산업용 컴퓨팅 자원을 필요로 하지 않고도 대규모 모델을 훈련할 수 있게 합니다.

저랭크 그래디언트 투영을 통한 메모리 효율성

GaLore는 옵티마이저가 처리하기 전에 그래디언트를 저차원 서브스페이스로 투영함으로써 메모리 사용량을 감소시킵니다. 이 방법은 딥 뉴럴 네트워크에서 그래디언트가 갖는 고유한 저랭크 구조를 활용하여 훈련 중 저장 및 조작해야 하는 데이터 양을 최소화합니다.

Adam과 같은 적응형 최적화 알고리즘에서는 옵티마이저 상태가 메모리 사용량의 큰 부분을 차지합니다. 이 투영을 적용함으로써 GaLore는 훈련 중 옵티마이저 상태 저장에 필요한 메모리를 82.5% 이상 감소시켰다고 보고되었습니다.

동적 서브스페이스 전환

전체 파라미터 학습 능력을 유지하고 모델이 파라미터 공간의 제한된 영역에 국한되지 않도록 하기 위해 GaLore는 동적 서브스페이스 전환 메커니즘을 활용합니다. 이 메커니즘은 훈련 과정에서 모델이 다양한 저랭크 서브스페이스를 탐색하도록 합니다.

이 전환들의 빈도는 그래디언트의 변화하는 저랭크 구조에 적응하면서 일관된 최적화 경로를 유지하도록 균형을 맞춥니다. 이를 통해 메모리 효율성과 최적화 성능 사이의 트레이드오프를 세밀하게 제어할 수 있습니다.

8비트 옵티마이저와의 통합

GaLore를 8비트 정밀도 옵티마이저와 결합하면 옵티마이저 상태를 양자화하여 메모리 효율성을 더욱 극대화합니다. 이러한 시너지 효과는 동일한 하드웨어 제약 내에서 모델 정확도나 수렴 속도를 손상시키지 않으면서 더 큰 모델을 훈련하거나 더 큰 배치 크기를 사용할 수 있게 합니다.

GaLore와 함께하는 8비트 최적화의 알고리즘 과정

1. Gradient Projection: 전체 정밀도 그래디언트를 투영 행렬을 사용해 저랭크 서브스페이스로 투영한 뒤 8비트 형식으로 양자화합니다.
2. Quantization: 투영된 그래디언트, 모델 가중치 및 옵티마이저 상태(예: Adam의 이동 평균)를 32비트 부동소수점에서 8비트 정수 표현으로 양자화합니다.
3. Optimizer Update: 8비트 양자화된 그래디언트를 업데이트합니다; 이 과정은 그래디언트를 부동소수점으로 역양자화하고, 업데이트 규칙을 적용한 뒤, 업데이트된 옵티마이저 상태를 다시 8비트로 양자화하는 것을 포함합니다.
4. De-quantization and Weight Update: 가중치를 처리하기 위해 부동소수점으로 역양자화합니다. GaLore는 이후 최종 투영을 사용해 역양자화된 저랭크 업데이트를 원래 파라미터 공간으로 다시 매핑한 뒤 가중치 업데이트를 적용합니다.

Hugging Face Transformers와의 구현

GaLore는 Hugging Face transformers 라이브러리(버전 4.39.0 이상)와 galore-torch 라이브러리에 통합되어 있습니다. 사용자는 TrainingArguments에서 옵티마이저를 galore_adamw, galore_adamw_8bit, galore_adafactor와 같은 옵션으로 지정하고 optim_target_modules를 통해 대상 모듈을 정의함으로써 GaLore를 구현할 수 있습니다.

레이어별 업데이트

메모리 사용량을 더욱 줄이기 위해 GaLore는 레이어별 가중치 업데이트를 지원합니다. 역전파 후 모든 레이어에 대해 한 번에 업데이트하는 대신, 옵티마이저는 PyTorch 포스트-누적 훅을 사용해 레이어당 하나씩 가중치를 업데이트합니다. 이 기능은 옵티마이저 이름에 _layerwise를 추가하여 활성화됩니다(예: galore_adamw_layerwise).