Google DeepMind의 Gemma 모델을 Hugging Face에서 LoRA와 QLoRA를 활용해 fine-tuning하여 메모리 효율적인 파라미터 최적화 달성

Context

대규모 언어 모델의 전체 가중치 fine-tuning은 메모리와 연산량이 많이 필요해 Colab이나 Kaggle 같은 무료 GPU 리소스나 엔터프라이즈 환경에서 비용 최적화가 어렵다. 개별 도메인에 맞게 모델을 적응시키는 데 드는 연산 비용이 중요한 최적화 지표가 되고 있다.

Technical Solution

• LoRA(Low-Rank Adaptation) 도입: 원본 모델은 동결하고 저차원 행렬로 분해된 어댑터 레이어만 학습하여 fine-tuning 파라미터 수 감소
• QLoRA 적용: BitsAndBytesConfig를 통해 기본 모델을 4-bit 정밀도로 양자화하여 메모리 효율성 극대화
• 타겟 모듈 지정: q_proj, o_proj, k_proj, v_proj, gate_proj, up_proj, down_proj 7개 레이어를 LoRA 어댑터 적용 대상으로 선택
• PyTorch/XLA 최적화: Gemma 모델을 PyTorch와 PyTorch/XLA 모두에서 실행 가능하게 구현하여 GPU와 TPU 사용자 모두 지원
• FSDP via SPMD 통합: Hugging Face transformers에서 PyTorch/XLA의 FSDP 구현을 지원하여 TPU 가속화를 통한 fine-tuning 속도 향상

Key Takeaway

PEFT(Parameter-Efficient Fine-Tuning) 기법, 특히 LoRA와 QLoRA 조합을 사용하면 전체 모델 크기를 동결 상태로 유지하면서 제한된 리소스 환경에서도 대규모 언어 모델의 도메인 특화 fine-tuning이 가능하다. 개발자는 타겟 모듈과 저차수 파라미터 설정을 통해 메모리-성능 트레이드오프를 직접 제어할 수 있다.