LoGRA, 저랭크 그래디언트 스케치로 LLM 강화학습 메모리 절감

배경

대규모 언어 모델의 추론 능력을 높이는 방법으로 강화학습 기반 후훈련이 널리 주목받고 있습니다. 그러나 강화학습은 일반적인 지도 미세조정보다 높은 메모리 비용을 요구합니다. 모델 가중치뿐 아니라 그래디언트, 옵티마이저 상태, 정책 업데이트와 동기화에 필요한 정보까지 처리해야 하기 때문입니다. 모델 규모가 커질수록 밀집 Adam과 같은 기존 방식은 GPU 메모리 한계에 빠르게 도달할 수 있습니다. LoGRA는 이 문제를 그래디언트 신호를 저장하는 방식에서 접근합니다.

LoGRA의 핵심 아이디어

LoGRA는 강화학습 업데이트에 필요한 유용한 그래디언트를 완전한 밀집 형태로 유지하는 대신, 저랭크 그래디언트 스케치로 표현합니다. 저랭크 표현은 원래 정보보다 작기 때문에 업데이트 과정의 메모리 사용량을 낮출 수 있습니다. 또한 압축된 표현을 모델 업데이트뿐 아니라 정책 동기화에도 활용해 관련 비용을 줄입니다.

하지만 그래디언트 압축만 적용하면 다른 문제가 생길 수 있습니다. 근사 그래디언트로 인해 업데이트 폭이 지나치게 커지면 정책이 갑자기 변하고, 학습 안정성이 흔들릴 수 있습니다. LoGRA는 이를 보완하기 위해 예측 KL 스텝 제어를 도입합니다. 실제 업데이트를 적용하기 전에 정책 변화량을 추정하고, 그 결과에 따라 업데이트 크기를 조절하는 방식입니다. 즉, 메모리 절감과 정책 변화 제한을 함께 설계합니다.

논문 요약에서 제시한 주요 결과는 다음과 같습니다.

• 추론 과제에서 성능 저하 없이 평균 학습 메모리를 최대 45.7% 절감했다;
• 밀집 Adam이 메모리 부족을 일으키는 조건에서도 단일 8-GPU 노드에서 27B 파라미터 모델을 1,100스텝 이상 안정적으로 학습했다;
• 구현은 Molt 라이브러리의 LoGRA 예제 스크립트로 공개됐다.

의미와 남은 과제

LoGRA의 의미는 단순히 GPU 메모리를 조금 아끼는 데 있지 않습니다. 기존 하드웨어에서는 실행하기 어려웠던 대규모 RL 실험을 더 작은 GPU 구성에서 시도할 가능성을 열어줍니다. 동시에 그래디언트 압축의 효과는 압축률만으로 판단할 수 없으며, 압축된 업데이트가 정책을 얼마나 안정적으로 변화시키는지도 함께 관리해야 한다는 점을 보여줍니다.

다만 현재 제공된 자료는 초록 중심이므로, 저랭크 차원별 비교, 과제별 세부 결과, 통신 비용, 다른 강화학습 알고리즘에서의 성능은 확인할 수 없습니다. 따라서 45.7%라는 수치는 모든 환경에서 보장되는 절감폭이 아니라 논문이 보고한 최대치로 이해해야 합니다. 다양한 모델 구조와 더 긴 학습 과정에서의 일반화 여부는 본문과 공개 코드에 대한 추가 검증이 필요합니다.

출처: Hugging Face Daily Papers