LLM 강화학습의 메모리 병목을 저랭크 그래디언트로 풀다
원제 LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches
추천 14댓글 2
Key Point
LLM 강화학습 훈련의 메모리 효율을 45% 이상 개선해 더 큰 모델과 리소스 제약 환경에서의 RL 훈련을 실용적으로 만든다.
핵심 요약
- 강화학습(RL)은 대규모언어모델(LLM)의 성능을 높였지만 메모리 수요가 채택의 걸림돌이 되고 있다.
- LoGRA는 유용한 학습 신호를 저랭크 그래디언트 스케치로 압축해 메모리 사용량을 줄이는 RL 후훈련 기법이다.
- 이 컴팩트한 표현은 모델 업데이트와 효율적인 정책 동기화 모두를 지원한다.
- 과도한 업데이트가 학습을 방해하지 않도록 predicted-KL 스텝 제어를 추가해 정책 변화를 사전 예측한 뒤 업데이트 크기를 조절한다.
- 추론 작업 전반에서 LoGRA는 성능 저하 없이 평균 훈련 메모리를 최대 45.7% 감축했다.
- 27B 파라미터 모델을 단일 8GPU 노드에서 1,100스텝 이상 안정적으로 훈련할 수 있게 하며, 이는 기존 Adam이 메모리 부족으로 실패하는 규모다.
- 코드는 Molt 라이브러리에서 이용 가능하다.