엔비디아, 강화학습에 추론 디코딩 도입해 속도 최대 1.8배 향상
NVIDIA 연구팀이 언어 모델의 강화학습 훈련 루프에 추측적 디코딩(Speculative Decoding)을 통합하여, 모델 출력 분포 손실 없이 롤아웃 생성 속도를 1.8배 향상시키는 기술을 선보였습니다. 이 방식은 기존의 처리량 증가를 위해 훈련 충실도를 희생해야 했던 문제를 해결하여, 모델 성능 저하 없이 완벽한 속도 개선을 이뤄냈다는 점에서 매우 중요합니다.
인공지능(AI) 인프라 기술 뉴스, AI 논문 요약 및 기술 동향, 대형 언어 모델(LLM) 및 머신러닝 섹션의 에디터 추천 글입니다.
수학적 추론, 코드 생성 또는 검증 가능한 작업을 위해 언어 모델에 대한 강화학습(RL) 사후 훈련(post-training)을 실행해 본 적이 있다면, GPU 클러스터가 롤아웃 생성 작업을 처리하는 동안 느리게 진행되는 진행 표시줄을 응시해 본 경험이 있을 것입니다. NVIDIA 연구팀은 추측적 디코딩(Speculative Decoding)을 RL 훈련 루프 자체에 통합하여 이 문제를 정확하게 해결하는 방법을 제안했으며, 타겟 모델의 정확한 출력 분포를 유지하는 방식으로 이를 구현했습니다. 연구팀은 vLLM 백엔드를 사용하는 NeMo RL v0.6.0에 추측적 디코딩을 직접 통합하여, 8B 모델 규모와 예상되는 235B 모델 규모 모두에서 무손실 롤아웃 가속을 구현했습니다. 최신 NeMo RL v0.6.0 릴리스에는 SGLang 백엔드, Muon 옵티마이저 및 YaRN 장문맥(long-context) 훈련과 함께 추측적 디코딩이 공식 지원 기능으로 포함되어 제공됩니다.
왜 롤아웃 생성이 병목 현상인가?
이 문제를 이해하려면 동기식 RL 훈련 단계가 어떻게 세분화되는지 아는 것이 도움이 됩니다. NeMo RL에서 각 단계는 데이터 로딩, 가중치 동기화 및 백엔드 준비(prepare), 롤아웃 생성(gen), 로그 확률 재계산(logprob), 정책 최적화(train)의 다섯 단계로 구성됩니다. 연구팀은 두 가지 워크로드 하에서 Qwen3-8B 모델을 사용하여 이를 측정했습니다. 추론 기능이 있는 모델의 훈련을 계속하는 'RL-Think'와 기본 모델에서 시작하여 처음부터 추론을 학습하는 'RL-Zero'입니다. 두 경우 모두 롤아웃 생성이 전체 단계 시간의 6572%를 차지했습니다. 로그 확률 재계산과 훈련은 합쳐봐야 약 2733% 정도만 차지합니다. 이로 인해 생성 단계가 가속화의 대상이 되는 유일한 단계이자, 롤아웃 측면 최적화의 한계를 결정하는 요소가 됩니다.
추측적 디코딩이 실제로 하는 일은 무엇인가? 추측적 디코딩은 더 작고 빠른 초안 모델(draft model)이 한 번에 여러 토큰을 제안하고, 더 큰 타겟 모델(실제로 훈련 중인 모델)이 기각 샘플링(rejection sampling) 절차를 사용하여 이를 검증하는 기술입니다. 이 기술의 핵심 속성이자 RL에서 중요한 이유는, 기각 절차가 타겟 모델이 자기회귀적(autoregressively)으로 해당 토큰을 생성한 것과 수학적으로 동일한 출력 분포를 생성한다는 것이 보장된다는 점입니다. 분포 불일치가 없고, 오프-정책(off-policy) 수정이 필요하지 않으며, 훈련 신호에도 변화가 없습니다. 이는 RL 사후 훈련에서 훈련 보상이 정책 자체의 샘플에 의존하기 때문에 매우 중요합니다. 비동기 실행, 오프-정책 재생 또는 저정밀도 롤아웃과 같은 방법은 모두 처리량을 높이기 위해 어느 정도의 훈련 충실도를 희생합니다. 반면 추측적 디코딩은 아무것도 희생하지 않습니다. 롤아웃은 타겟 모델이 자체적으로 생성했을 것과 분포적으로 완벽히 동일하지만, 단지 더 빠르게 생성될 뿐입니다.
시스템 통합의 과제 서빙 백엔드에 초안 모델을 추가하는 것은 간단합니다. 하지만 RL 훈련 루프에 추가하는 것은 그렇지 않습니다. 정책이 업데이트될 때마다 롤아웃 엔진은 새로운 가중치를 받아야 합니다. 또한 초안 모델은 계속해서 진화하는 정책과 정렬을 유지해야 합니다. 로그 확률, KL 페널티 및 GRPO 정책 손실은 모두 초안 모델이 아닌 타겟(검증자) 정책을 기준으로 계산되어야 합니다. 그렇지 않으면 최적화 대상이 알게 모르게 손상됩니다. NVIDIA 연구팀은 NeMo RL에서 '듀얼 패스 아키텍처(two-path architecture)'로 이 문제를 해결합니다. 일반 패스는 기본적인 다중 토큰 예측(MTP) 지원 없이도 사전 훈련된 모델과 함께 작동하는 초안 프레임워크인 EAGLE-3를 사용합니다. 또한 기본 제공 MTP 헤드가 포함된 모델을 위한 네이티브 패스도 제공됩니다. 온라인 초안 적응이 활성화되면, MegatronLM 검증자의 순전파(forward pass)에서 얻은 은닉 상태(hidden states)와 로그 확률이 캐시되어 기울기가 분리된 경로(gradient-detached pathway)를 통해 초안 헤드를 감독하는 데 재사용되므로, 초안 훈련이 정책 기울기 신호를 방해하지 않습니다.
8B 규모에서의 측정 결과 32개의 GB200 GPU(8개의 GB200 NVL72 노드, 노드당 4개의 GPU) 환경에서 EAGLE-3는 생성 대기 시간을 단축시킵니다.