Hugging Face에서 Gemma 모델 파인튜닝

Hugging Face는 Parameter-Efficient Fine-Tuning(PEFT)을 사용하여 Google DeepMind의 Gemma 모델을 파인튜닝하는 상세 워크플로우를 소개했습니다. 이 접근 방식은 전체 가중치 학습에 비해 훨씬 낮은 메모리와 연산 요구량으로 2 억 및 7 억 파라미터 Gemma 모델을 특정 데이터셋에 맞게 조정할 수 있게 합니다.

Gemma를 위한 Low-Rank Adaptation (LoRA)

Low-Rank Adaptation(LoRA)은 Gemma 모델에 강조된 주요 PEFT 기법입니다. 모든 모델 파라미터를 업데이트하는 대신, LoRA는 원본 가중치를 고정하고 저랭크 행렬로 구성된 소수의 어댑터 레이어만 학습합니다. 이를 통해 파인튜닝 시 계산 오버헤드가 크게 감소합니다.

Gemma와 함께 LoRA를 구현하기 위해 Hugging Face PEFT 라이브러리는 사용자가 특정 선형 레이어를 타깃으로 지정할 수 있게 합니다. 제공된 예시에서는 다음 모듈들을 어댑테이션 대상으로 지정합니다:

  • q_proj
  • o_proj
  • k_proj
  • v_proj
  • gate_proj
  • up_proj
  • down_proj

QLoRA를 통한 메모리 최적화

메모리 사용량을 더욱 줄이기 위해 Hugging Face는 QLoRA 사용을 권장합니다. 이 방법은 bitsandbytes 라이브러리를 이용해 기본 모델을 4‑bit 정밀도로 양자화합니다. from_pretrained 메서드에 BitsAndBytesConfig를 전달하면, nf4 양자화 타입과 bfloat16 연산 dtype 등을 사용해 Gemma를 메모리 효율적인 형식으로 로드할 수 있어, 일반 소비자용 GPU나 Google Colab 같은 무료 플랫폼에서도 파인튜닝이 가능해집니다.

구현 워크플로우: 인용문 생성 학습

trl 라이브러리의 SFTTrainer를 활용해 Hugging Face는 Gemma‑2b를 특정 형식(인용문 뒤에 저자)으로 인용문을 생성하도록 파인튜닝하는 실용적인 예시를 보여줍니다.

프로세스는 다음과 같습니다:

  1. 모델 로딩: BitsAndBytesConfig를 통해 4‑bit 양자화된 모델을 로드합니다.
  2. 데이터셋 준비: Abirate/english_quotes 데이터셋을 사용하고 모델에 맞게 토크나이징합니다.
  3. 구성: 랭크(r)가 8인 LoraConfig를 적용합니다.
  4. 학습: 학습률 2e-4와 paged_adamw_8bit 옵티마이저로 트레이너를 실행합니다.

하드웨어 가속: TPU에서 PyTorch/XLA와 FSDP

Hugging Face transformers 라이브러리의 Gemma 모델은 PyTorch와 PyTorch/XLA 모두에 최적화되어 있어 GPU와 Cloud TPU 모두에 배포할 수 있습니다.

TPU 사용자를 위해 Hugging Face는 SPMD(단일 프로그램, 다중 데이터)를 통한 Fully Sharded Data Parallel(FSDP) 경험을 개선했습니다. transformers.Trainer에 fsdp_config를 추가하고 GemmaDecoderLayer를 래핑하며 xla_fsdp_v2를 활성화하면, TPU 하드웨어에서 파인튜닝 속도가 크게 가속됩니다.

TPU 구성 요약

기능 설정
래핑할 레이어 GemmaDecoderLayer
FSDP 모드 full_shard
XLA FSDP v2 활성화 (True)
XLA FSDP Grad 체크포인팅 활성화 (True)

Sources