시스템 통합 투기적 디코딩을 통한 강화학습 사후 학습 롤아웃 가속화

LLM의 추론 능력을 높이는 강화학습(RL) 과정에서 가장 큰 병목인 롤아웃 생성 단계를 투기적 디코딩 기법으로 해결했다. 기존의 효율화 방식과 달리 모델의 출력 분포를 전혀 해치지 않으면서도 훈련 시간을 획기적으로 단축할 수 있는 시스템적 기반을 마련했다.

강화학습 사후 학습은 모델이 직접 생성한 샘플(Rollout)을 바탕으로 정책을 업데이트하는데, 이 생성 과정이 전체 훈련 시간의 65~72%를 차지하는 심각한 병목 현상을 일으킨다. 기존에는 이를 해결하기 위해 낮은 정밀도로 생성하거나 과거 데이터를 재사용하는 방식을 썼으나, 이는 모델의 원래 출력 분포와 차이를 발생시켜 학습의 질을 떨어뜨릴 위험이 있다.

이 논문은 투기적 디코딩(Speculative Decoding)을 해결책으로 제시한다. 이는 가벼운 드래프트 모델이 여러 토큰을 미리 제안하고 타겟 모델이 이를 한 번에 검증하는 방식이다. 타겟 모델의 거부 샘플링(Rejection Sampling) 절차를 거치기 때문에 최종 결과물은 타겟 모델이 직접 생성한 것과 수학적으로 동일한 분포를 유지한다.

결과적으로 학습 알고리즘 입장에서는 데이터의 품질 변화 없이 데이터 공급 속도만 빨라지는 효과를 얻는다. 특히 모델이 업데이트될 때마다 드래프트 모델을 실시간으로 정렬하는 시스템적 설계를 통해 학습 전 과정에서 가속 성능을 유지하도록 구현했다.

NeMo-RL 프레임워크 내부에 vLLM 서빙 백엔드를 통합하여 투기적 디코딩 롤아웃 엔진을 구축했다. 학습자가 매 단계 정책 가중치를 업데이트하면 롤아웃 엔진은 이를 수신하여 타겟 모델과 드래프트 모델의 동기화를 수행한다.

드래프트 메커니즘으로 EAGLE-3 기반의 일반 드래프팅 경로와 모델 내장형 MTP(Multi-Token Prediction) 헤드를 사용하는 네이티브 경로를 모두 지원한다. EAGLE-3를 사용할 때는 타겟 모델의 Forward Pass에서 생성된 Hidden State를 캐싱하여 드래프트 모델의 학습 데이터로 재사용함으로써 추가적인 연산 오버헤드를 최소화했다.

가속 성능 분석을 위해 Amdahl의 법칙을 변형한 수식을 적용했다. 전체 단계 시간 T_step 대비 생성 시간 T_gen의 비율과 평균 수락 길이 α를 기반으로 가속 상한을 계산한다. [T_gen / T_step → 생성 비중 계산 → 1 / ( (T_gen/α) + (1-T_gen) ) → 이론적 가속 배수 도출]

8B 규모의 Qwen3 모델을 사용한 실험에서 EAGLE-3 드래프팅은 RL-Zero 설정에서 생성 지연시간을 100.0초에서 56.6초로 줄여 1.8배의 가속을 달성했다. 전체 RL 단계 시간 기준으로는 1.41배의 속도 향상을 기록했다.

n-gram 기반의 모델 프리 드래프팅과 비교했을 때, n-gram은 2.47의 수락 길이를 기록했음에도 불구하고 검증 오버헤드로 인해 오히려 기본 오토레그레시브 방식보다 0.7배 느려지는 결과를 보였다. 이는 단순한 수락 길이보다 시스템적 효율성이 중요함을 시사한다.

시뮬레이션 결과, 235B 규모 모델을 2048개의 GB200 GPU에서 실행할 경우 롤아웃 가속은 약 3.5배에 달하며, 이를 통해 전체 훈련 시간을 약 2.5배 단축할 수 있는 것으로 나타났다.

시스템 아키텍처는 MegatronLM(학습자)과 vLLM(롤아웃 엔진) 간의 긴밀한 결합으로 구성된다. 온라인 드래프트 적응(Online Draft Adaptation) 기능을 통해 학습 과정에서 변화하는 정책 분포에 드래프트 모델을 실시간으로 맞춘다. 이때 .detach() 경로를 사용하여 드래프트 모델의 손실 함수가 메인 정책의 그래디언트 신호를 간섭하지 않도록 설계했다.

비동기 RL 모드에서는 생성이 학습과 겹쳐서 진행되므로 임계 경로(Critical Path) 상의 생성 비중이 줄어든다. 이 경우 투기적 디코딩의 절대적인 가속 배수는 동기식보다 낮아질 수 있으나, 여전히 유의미한 엔드투엔드 가속을 제공하며 두 메커니즘이 상호 보완적으로 작동함을 확인했다.

시뮬레이션 분석에 따르면 모델 규모가 커질수록 생성 단계의 비중이 높아지기 때문에 투기적 디코딩을 통한 가속 잠재력이 더욱 커진다. 특히 GB200과 같은 최신 하드웨어 환경에서 FP8 정밀도를 사용할 때의 오버헤드와 이득을 정밀하게 계산하여 실전 배치 가이드를 제시했다.