ServiceNow Research가 LLM 강화학습에서 인플라이트 가중치 업데이트 도입으로 추론 처리량과 정책 기반 데이터 수집 간 트레이드오프 해결

Context

기존 LLM 강화학습 방식은 높은 추론 처리량과 정책 기반 데이터 수집 간 근본적인 트레이드오프를 가지고 있다. 추론 서버가 높은 처리량을 위해 대규모 배치로 여러 최적화 단계의 데이터를 생성하면, 각 최적화 단계마다 현재 정책과 데이터 수집에 사용된 정책 간 지연이 증가하여 데이터가 점진적으로 정책 기반에서 벗어나 학습 효율성이 저하된다.

Technical Solution

• 인플라이트 가중치 업데이트 도입: 옵티마이저 단계 완료 후마다 추론 서버의 가중치를 갱신하되 추론을 중단하지 않고 새 가중치 수신 시간만 일시 정지
• 단순화된 GRPO 알고리즘 적용: 신뢰 영역 중요도 가중치 클래핑, 긴 시퀀스 필터링, 보상 형성을 제거하고 배치 내 시퀀스 개수로 정규화
• 모듈식 아키텍처 설계: 추론 및 훈련 컴포넌트 간 명확한 인터페이스 계약으로 SGLang, vLLM, DeepSpeed, FSDP 등 다양한 인프라 통합 지원
• HTTP POST 기반 프로세스 그룹 초기화: 훈련자 0이 /init_process_group 요청으로 모든 추론 서버의 프로세스 그룹 초기화
• NCCL을 통한 가중치 브로드캐스트: 옵티마이저 단계 완료 후 /request_weight_update로 가중치 순서와 형태를 알린 후 NCCL로 가중치 전송

Impact

PipelineRL은 7B 모델에서 Open-Reasoner-Zero와 동등 이상의 성능을 AIME 2024 및 MATH 500 벤치마크에서 달성했다. 7B 모델은 2개 노드에서 약 3.5일, 32B 모델은 4개 노드에서 약 6일의 계산량으로 훈련되었다.

Key Takeaway

인플라이트 가중치 업데이트는 KV 캐시의 구식 키와 값이 존재함에도 불구하고 학습 안정성을 해치지 않으며, 이는 대규모 LLM 강화학습에서 처리량과 정책 기반성을 동시에 확보할 수 있음을 시사한다.