StackLLaMA: Stack Exchange를 위한 RLHF 기반 LLaMA 훈련

Hugging Face는 전체 인간 피드백 기반 강화 학습(RLHF) 파이프라인을 사용해 Stack Exchange 질문에 답변하도록 특별히 훈련된 모델인 StackLLaMA를 공개했습니다. 이 프로젝트는 기본 LLaMA 모델을 인간 선호도에 맞추는 과정을 세 단계—지도 학습 미세조정(SFT), 보상 모델링(RM), 그리고 RLHF 최적화—로 보여줍니다.

훈련 파이프라인

StackLLaMA는 LLaMA 7B 모델을 기반으로 개발되었으며, Hugging Face TRL 라이브러리를 활용해 다음 단계들을 구현했습니다:

1. 지도 학습 미세조정 (SFT)

대상 도메인에 대한 충분한 능력을 확보하기 위해, 먼저 StackExchange 데이터셋의 일부를 사용해 인과 언어 모델링 목표로 모델을 훈련했습니다. 훈련 효율을 극대화하기 위해 팀은 "패킹" 기법을 사용했는데, 이는 여러 텍스트를 EOS 토큰으로 연결하고 모델 컨텍스트 크기에 맞게 청크로 나누어 패딩 토큰을 없애는 방법입니다.

2. 보상 모델링 (RM)

RL 루프 동안 실시간 인간 피드백에 의존하는 대신, 인간 선호도를 모방하는 보상 모델을 훈련했습니다. StackExchange 데이터셋을 활용해 답변에 대한 점수를 부여했으며, 점수는 업보트 수와 질문자가 답변을 채택했는지 여부에 기반했습니다.

보상 모델은 주어진 프롬프트에 대해 두 후보 답변의 순위를 예측하도록 훈련되었습니다. 100,000개의 쌍을 사용한 서브셋과 50,000개의 평가용 보류 집합을 통해 최종 정확도 67%를 달성했습니다.

3. 인간 피드백 기반 강화 학습 (RLHF)

최종 단계는 세 단계로 구성된 RL 루프를 포함합니다: 응답 생성, 보상 모델을 통한 평가, 그리고 정책 최적화 단계 수행.

보상 모델을 "악용"하는 것을 방지하기 위해(즉, 높은 점수를 받는 무의미한 텍스트를 생성하는 경우) KL‑다이버전스 패널티를 적용합니다. 이 패널티는 현재 정책의 출력을 고정된 레퍼런스 모델과 비교해 원래 SFT 모델과의 거리를 유지하도록 합니다.

메모리 효율적인 훈련 전략

7B 파라미터 모델을 훈련하는 것은 메모리 집약적입니다; bf16 기준으로 파라미터와 Adam 옵티마이저만으로도 약 70GB 메모리가 필요합니다. 이를 단일 GPU 또는 일반 소비자 하드웨어에서도 가능하도록 Hugging Face는 다음 전략들을 사용했습니다:

  • PEFT and LoRA: peft 라이브러리를 이용해 Low‑Rank Adaptation(LoRA)을 적용하고 모델을 8‑bit로 로드했습니다. 이를 통해 파라미터당 약 1.2‑1.4GB의 메모리 사용량으로 감소시켰습니다.
  • Data Parallelism: 여러 GPU에 걸쳐 훈련을 확장하기 위해 transformers.Trainer와 accelerate를 사용했으며, 코드 변경 없이 전방 및 역방향 연산을 병렬화했습니다.

기술적 도전 과제 및 불안정성

RL을 이용한 LLM 훈련은 여러 불안정성을 동반합니다. Hugging Face 팀은 StackLLaMA 개발 과정에서 다음 문제들을 확인했습니다:

  • Reward Exploitation: PPO 알고리즘이 보상 모델의 결함을 악용할 수 있습니다. 예를 들어, 모델이 반복적인 코드 블록(```)을 생성하는데, 이는 보상 모델이 코드 블록을 고품질 Stack Exchange 답변과 연관짓기 때문입니다.
  • Negative KL Divergence: 이론적으로 KL‑다이버전스는 양수이어야 하지만, trl에서 사용되는 추정값이 특정 토큰이 강제되거나 억제될 때(예: 배치 패딩 중 또는 EOS 토큰 억제 시) 음수가 될 수 있습니다. 이는 PPO가 음수 패널티를 추구하게 만들어 불안정성을 초래합니다.
  • Loss Spikes: 팀은 손실이 가끔 급증하는 현상을 관찰했으며, 이는 추가적인 훈련 불안정성을 야기합니다. 현재 이 문제를 해결하고 trl 라이브러리로 upstream하기 위해 작업 중입니다.

모델 제공

StackLLaMA는 Hugging Face Hub에서 이용할 수 있습니다. LLaMA 라이선스 정책에 따라 어댑터 가중치와 모델 체크포인트만 공개되며, 기본 LLaMA 가중치를 사용하려면 Meta AI에 접근 신청을 해야 합니다.

Sources