메뉴
BL
MarkTechPost 29일 전

엔비디아, 강화학습에 추론 디코딩 도입해 속도 최대 1.8배 향상

IMP
8/10
핵심 요약

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는 생성 대기 시간을 단축시킵니다.

원문 보기
원문 보기 (영어)
Artificial Intelligence AI Infrastructure Tech News AI Paper Summary Technology AI Shorts Applications Editors Pick Language Model Large Language Model Machine Learning Staff If you have been running reinforcement learning (RL) post-training on a language model for math reasoning, code generation, or any verifiable task, you have almost certainly stared at a progress bar while your GPU cluster burns through rollout generation. A team of researchers from NVIDIA proposes a precise fix by integrating speculative decoding into the RL training loop itself, and do it in a way that preserves the target model's exact output distribution. The research team integrated speculative decoding directly into NeMo RL v0.6.0 with a vLLM backend, delivering lossless rollout acceleration at both 8B and projected 235B model scales.The latest NeMo RL v0.6.0 release officially ships speculative decoding as a supported feature alongside the SGLang backend, the Muon optimizer, and YaRN long-context training. Why Rollout Generation is the Bottleneck To understand the problem, it helps to know how a synchronous RL training step breaks down. In NeMo RL, each step consists of five stages : data loading, weight synchronization and backend preparation (prepare), rollout generation (gen), log-probability recomputation (logprob), and policy optimization (train). The research team measured this breakdown on Qwen3-8B under two workloads — RL-Think , which continues training a reasoning-capable model, and RL-Zero , which starts from a base model and learns reasoning from scratch. In both cases, rollout generation accounts for 65–72% of total step time. Log-probability recomputation and training together take only about 27–33%. This makes generation the only stage worth targeting for acceleration, and the one that determines the ceiling for any rollout-side optimization. What Speculative Decoding Actually Does Speculative decoding is a technique where a smaller, faster draft model proposes several tokens at once, and the larger target model (the one you are actually training) verifies them using a rejection sampling procedure. The key property and why it matters for RL, is that the rejection procedure is mathematically guaranteed to produce the same output distribution as if the target model had generated those tokens autoregressively. No distribution mismatch, no off-policy corrections needed, no change to the training signal. This is important because in RL post-training, the training reward depends on the policy's own samples. Methods like asynchronous execution, off-policy replay, or low-precision rollouts all trade some amount of training fidelity for throughput. Speculative decoding trades nothing: the rollouts are identical in distribution to what the target model would have generated on its own, just produced faster. The System Integration Challenge Adding a draft model to a serving backend is straightforward. Adding one to an RL training loop is not. Every time the policy updates, the rollout engine must receive new weights. The draft model must remain aligned with the evolving policy. Log-probabilities, KL penalties, and the GRPO policy loss must all be computed against the target (verifier) policy not the draft or the optimization target is silently corrupted. The NVIDIA research team handles this in NeMo RL with a two-path architecture . The general path uses EAGLE-3, a drafting framework that works with any pretrained model without requiring native multi-token prediction (MTP) support. A native path is also available for models that ship with built-in MTP heads. When online draft adaptation is enabled, the hidden states and log-probabilities from the MegatronLM verifier forward pass are cached and reused to supervise the draft head via a gradient-detached pathway, so draft training never interferes with the policy gradient signal. Measured Results at 8B Scale On 32 GB200 GPUs (8 GB200 NVL72 nodes, 4 GPUs per node), EAGLE-3 reduces generation latency from 100 seconds to 56.6 seconds on RL-Zero — a 1.8× generation speedup. On RL-Think, it drops from 133.6 seconds to 87.0 seconds, a 1.54× speedup. Because log-probability re-computation and training are unchanged, these generation-side gains translate to overall step speedups of 1.41× on RL-Zero and 1.35× on RL-Think. Validation accuracy on AIME-2024 evolves identically under autoregressive and speculative decoding throughout training, confirming that the lossless guarantee holds in practice. The research team also tests n-gram drafting as a model-free speculative baseline. Despite achieving acceptance lengths of 2.47 on RL-Zero and 2.05 on RL-Think, n-gram drafting is slower than the autoregressive baseline in both settings — 0.7× and 0.5× respectively. This is a critical finding for practitioners: a positive acceptance length is necessary but not sufficient. If the verification overhead is high enough, speculation makes things worse. Three Configuration Decisions That Determine Realized Speedup The research team isolates three operational choices that practitioners must get right. Draft initialization matters more than generic drafting ability. An EAGLE-3 draft initialized on the DAPO post-training dataset achieves a 1.77× generation speedup on RL-Zero, while a draft initialized on the general-purpose UltraChat and Magpie datasets achieves only 1.51× at the same draft length. The draft must be aligned with the actual rollout distribution encountered during RL, not just a broad chat distribution. Draft length has a non-obvious optimum. At draft length k=3, RL-Zero achieves 1.77× speedup and RL-Think achieves 1.53×. Increasing to k=5 raises the acceptance length but drops speedup to 1.44× on RL-Zero and 0.84× on RL-Think — the latter already slower than autoregressive. At k=7, RL-Zero drops further to 1.21× and RL-Think to 0.71×. The contrast matters: RL-Zero's rollouts are generated from a base model starting with short outputs, making them easier for the draft to predict even at high k. RL-Think's fully developed reasoning traces are harder to speculate over, so the overhead of longer drafts erases the benefit sooner. More speculative work per step can erase the benefit of higher acceptance entirely, especially in harder generation regimes. Online draft adaptation — updating the draft during RL using rollouts generated by the current policy helps most when the draft is weakly initialized. For a DAPO-initialized draft, offline and online configurations perform nearly identically (1.77× vs. 1.78× on RL-Zero). For a UltraChat-initialized draft, online updating improves speedup from 1.51× to 1.63× on RL-Zero. Interaction with asynchronous execution was also tested directly at 8B scale not just in simulation. The research team ran RL-Think at policy lag 1 in a 16-node non-colocated configuration, with 12 nodes dedicated to generation and 4 to training. In asynchronous mode, most of rollout generation is already hidden behind log-probability re-computation and policy updates, so the relevant quantity is the exposed generation time that remains on the critical path. Speculative decoding reduces that exposed generation time from 10.4 seconds to 0.6 seconds per step and lowers effective step time from 75.0 seconds to 60.5 seconds (1.24×). The gain is smaller than in synchronous RL — expected, since asynchronous overlap already hides much of the rollout cost — but it confirms that the two mechanisms are genuinely complementary rather than redundant. Projected Gains at 235B Scale Using a proprietary GPU performance simulator calibrated to device-level compute, memory, and interconnect characteristics, the research team projected speculative decoding gains at larger scales. For Qwen3-235B-A22B running synchronous RL on 512 GB200 GPUs, draft length k=3 with an acceptance length of 3 tokens yields a 2.72× rollout speedup and a 1.70× end-to-end speedup. At the most favorable simulated operating point — Qwen3-235B-A22B on 2048 GB200 GPUs w