24GB 소비자 GPU에서 RLHF를 이용한 20B LLM 파인튜닝

TL;DR

Hugging Face는 trl 라이브러리가 이제 peft와 8‑bit 양자화와 함께 작동한다는 것을 발표했으며, 이를 통해 24 GB 소비자 GPU 하나에서 20 B 파라미터 LLM의 RLHF 파인튜닝이 가능해졌습니다. 이는 다중 GPU 모델 병렬 설정 없이도 대규모 RL 파인튜닝을 저렴하고 접근 가능하게 만듭니다.

RLHF가 효율적인 파인튜닝을 필요로 하는 이유

인간 피드백을 이용한 강화 학습(RLHF)은 일반적으로 세 단계로 진행됩니다: (1) 지시문에 대한 감독 파인튜닝, (2) 인간 주석으로부터 보상 모델을 학습, (3) 보상 모델을 사용한 PPO 기반 RL 파인튜닝. RL 단계에서는 각 GPU에 모델 두 개(활성 모델과 기준 모델)가 필요해 10 B 파라미터를 초과하는 모델은 단일 장치 메모리를 초과합니다.

TRL: PPO 기반 RL을 위한 라이브러리

trl은 언어 모델의 PPO 훈련을 위한 고수준 API를 제공합니다. 🤗 Accelerate를 활용해 단일 장치 또는 분산 설정에서 실행할 수 있습니다. PPO 루프는 KL‑정규화 보상을 계산하기 위해 활성 모델(업데이트 중)과 기준 모델(동결) 두 개가 필요해 메모리 사용량이 두 배가 됩니다.

PEFT와 8‑Bit 양자화를 통한 메모리 사용량 감소

8‑Bit 행렬 곱셈

  • 8‑bit 양자화(LLM.int8())는 파라미터당 1바이트로 가중치를 저장해 float32 대비 모델 크기를 4배 줄입니다.
  • 이 방법은 각 선형 레이어를 이상치 처리용 float16 부분과 대량 int8 부분으로 분할하여 정확도를 유지하면서 속도를 높입니다.

PEFT를 통한 저랭크 적응(LoRA)

  • LoRA는 사전 학습된 가중치를 고정하고 어텐션 블록의 query와 value 투영에 저랭크 행렬(A와 B)을 삽입합니다.
  • 학습 가능한 파라미터는 어댑터 파라미터만으로, 옵티마이저 메모리를 크게 줄입니다.
  • 전방 및 역전파는 추가 행렬 곱셈 때문에 대략 두 배 느리지만, 메모리 절감으로 소비자 하드웨어에서 20 B 모델 훈련이 가능합니다.

24 GB GPU에서 20 B 모델을 위한 엔드‑투‑엔드 파이프라인

단계 1 – 8‑Bit 정밀도로 모델 로드

model = AutoModelForCausalLM.from_pretrained(
    "EleutherAI/gpt-neox-20b",
    load_in_8bit=True,
    device_map="auto",
)

8‑bit 로드 시 메모리가 ~80 GB(float32)에서 ~20 GB로 감소해 24 GB 카드에 여유롭게 맞습니다.

단계 2 – PEFT로 학습 가능한 LoRA 어댑터 연결

from peft import get_peft_model, LoraConfig
config = LoraConfig(r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], bias="none")
model = get_peft_model(model, config)

옵티마이저 상태에 저랭크 행렬만 저장되므로 옵티마이저 메모리가 수 기가바이트에서 수백 메가바이트 수준으로 감소합니다.

단계 3 – 기준 및 활성 로짓에 단일 모델 사용

PEFT의 disable_adapters 컨텍스트 매니저는 LoRA 레이어를 일시적으로 비활성화해 동일한 기본 모델이 기준 로짓을 생성하도록 합니다:

with model.disable_adapter():
    ref_logits = model(input_ids)
# active logits are computed with adapters enabled
active_logits = model(input_ids)

두 번째 전체 모델 복사가 필요 없어 메모리 사용량을 추가로 줄입니다.

훈련 스크립트 개요

  1. clm_finetune_peft_imdb.py – IMDB 감성 데이터셋(한 epoch)에서 LoRA 어댑터를 이용한 인과 언어 모델 파인튜닝.
  2. merge_peft_adapter.py – 추론 또는 추가 훈련을 위해 LoRA 가중치를 기본 모델에 병합.
  3. gpt-neo-20b_sentiment_peft.py – IMDB 감성 분류기를 보상 모델로 사용해 PPO 파인튜닝으로 긍정적인 영화 리뷰를 생성.

모든 스크립트는 NVIDIA RTX 4090(24 GB)에서 실행되었습니다. 전체 훈련은 🤗 연구 클러스터의 단일 A100에서도 테스트되었습니다.

결과

  • 손실 곡선은 IMDB에서 감독 LoRA 파인튜닝 1 epoch 후 안정적인 수렴을 보여줍니다.
  • PPO 동안 평균 보상이 꾸준히 증가해 모델이 더 긍정적인 리뷰를 생성함을 나타냅니다.
  • 전체 파이프라인이 단일 24 GB GPU에서 실행되어 RLHF가 더 이상 다중 GPU 클러스터에 국한되지 않음을 입증했습니다.

커뮤니티에 미치는 영향

  • 진입 장벽 감소 – 연구자와 개발자는 소비자 등급 하드웨어에서 RLHF를 실험할 수 있습니다.
  • 오픈소스 재현성 – 모든 코드와 어댑터가 Hugging Face Hub에 호스팅되어 파인튜닝된 산출물을 쉽게 공유할 수 있습니다.
  • 확장 가능한 기반 – 동일한 접근법을 데이터 병렬을 통해 더 큰 모델에도 적용할 수 있으며, 다중 GPU 지원이 추가되면 확장 가능합니다.

열린 질문 및 향후 작업

  • 다중 GPU 스케일링 – 여러 GPU에 걸친 데이터 병렬과의 통합은 얼마나 잘 작동할까요?
  • 훈련 속도 – LoRA는 오버헤드를 추가합니다; 더 빠른 커널이나 혼합 정밀도 전략을 탐색하면 완화될 수 있습니다.
  • 다양한 RL 알고리즘 – PPO가 기본이지만, 다른 RL 방법(예: DPO) 통합은 적용 범위를 넓힐 수 있습니다.

참고 문헌

Sources