24GB 일반 GPU에서 RLHF로 20B LLM 미세 조정하기

요즘 허깅페이스에서 개발하는 라이브러리와 작성하는 블로그 글들이 굉장히 트렌디하고 내용이 알차네요. 👍 본문 링크(https://huggingface.co/blog/trl-peft)입니다. ✔️ 수많은 ChatGPT 분석 글처럼, LLM을 RLHF으로 훈련하는 방법은 ① 명령 프롬프트와 레이블로 LLM 미세 조정 ② 출력 결과 순위 정해주는 보상 모델 훈련 ③ RL(PPO) 통해 프롬프트와 보상 모델로 LLM 추가 미세 조정하는 순. ✔️ 현재까지 '명령으로 미세 조정된 버전의 오픈소스 최고 LLM'은 BLOOMZ, Flan-T5, Flan-UL2, OPT-IML 정도. (LLaMA는 뛰어난 잠재력이 있지만 아직 천연의 아이인 셈이라고요, ChatGPT 같지 않다고 욕하지 마시길! 😌) ✔️ 허깅페이스 TRL = RL로 LM 미세 조정하는 걸 지원하는 라이브러리. 근데 RL 미세 조정 시에는 액티브 모델이 원래 동작 또는 분포에서 많이 벗어나는 걸 막기 위해 레퍼런스 모델 로짓 값 필요함, 결국 이는 GPU 당 모델 사본이 2개 필요하다는 얘기, 메모리 문제. ✔️파라미터 1B 당 float32 정밀도 = GPU 메모리 4GB, float16 = 2GB, int8 = 1GB, AdamW 옵티마이저 사용할 경우 8GB 필요 (공식처럼 외워라! 😌) ✔️ 대형 모델 훈련을 위해서는 당연히 병렬 처리 패러다임 적용이 필요, 데이터 병렬의 경우 TRL이 지원해 주는데 모델 병렬은 통신 프로토콜 등의 문제로 Megatron-DeepSpeed, Nemo 같은 프레임워크 별도 필요함. ✔️ 일단 데이터 병렬 처리에 국한하여 효율화해 볼까? (1) 8비트 행렬 곱셈: 은닉 상태 값의 이상치는 float16, 정상 값은 int8로 연산 수행하자. (2) LoRA: 사전 훈련 가중치 동결하고 Q, V 어텐션 행렬의 낮은 랭크 버전을 생성하여 미세 조정하자. 이로써 적은 파라미터 학습, 적은 GPU 메모리 소모, 대신에 순방향, 역방향 연산이 2배 정도 느림. ✔️ 허깅페이스 PEFT = LLM의 어댑터 생성과 미세 조정을 지원하는 라이브러리. ✔️ 24GB 단일 GPU에서 20B 파라미터 모델을 미세 조정하려면? ① 액티브 모델을 int8로 로딩 ② PEFT 이용해서 액티브 모델에 어댑터 (LoRA) 추가 ③ 필요시 어댑터 끌 수 있으므로 모델 사본 2개 필요 없이 같은 모델로 PPO에 대한 액티브 및 레퍼런스 모델 로짓을 구할 수 있음. ✔️ 예시로 GPT-NeoX (20B)에 대해 IMDB 영화 리뷰와 긍정 및 부정 레이블 데이터셋으로 LoRA 미세 조정 → 적응된 가중치 결합하고 허브에 저장 → LoRA 다시 추가하여 RL로 미세 조정 (이때, IMDB 감성 분류 모델을 보상 모델로 사용) (https://huggingface.co/docs/trl/main/en/sentiment_tuning_peft)