Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training

논문 링크

대규모·장문맥 RL 사후학습에서 Speculative Decoding을 실용화하는 두 가지 설계

TL;DR : RL 사후학습의 벽-클록 시간을 지배하는 롤아웃 생성을 speculative decoding으로 가속하고, 여기에 온라인 드래프트 공동학습(co-training) 을 더해 정책이 진화해도 드래프트가 뒤처지지 않게 한다. 다만 기존 시스템은 고급 드래프트가 쓰는 branch attention(문맥 병렬화 미지원) 과 여러 스테이지에 흩어진 target feature(파이프라인 병렬화 문제) 를 다루지 못했다. 이 논문은 (1) branch attention을 causal 주시퀀스 + rank-local branch로 분해해 병합하는 CP 기법과, (2) 파이프라인 스케줄 밖에서 target feature를 전달하는 TapChannel 을 제안해, 8B부터 122B까지 1.16–1.88× 의 end-to-end 가속을 달성했다(근거: §3.3, Tab. 1).

핵심 아이디어

RL 사후학습에서 speculative decoding(SD)은 작은 드래프트가 여러 토큰을 먼저 생성하고 큰 정책 모델이 병렬로 검증(verification)함으로써, 출력 분포를 바꾸지 않고 자기회귀 생성을 가속한다(근거: §1). 드래프트를 고정 하면 정책이 진화할수록 수용 길이(acceptance length)가 떨어지므로, 드래프트를 정책의 롤아웃 토큰으로 온라인 공동학습 하면 가속 폭이 더 커진다(근거: §1). 문제는 EAGLE-3, DFlash, DSpark 같은 최신 드래프트가 (a) branch 구조의 어텐션 을 쓰고 (b) 타깃 모델의 중간 hidden state 를 소비한다는 점이다(근거: §1, §2.2). 이 둘은 각각 표준 causal context parallelism(CP)과 pipeline parallelism(PP)의 한계와 정면으로 충돌한다. 저자들의 핵심 주장은 간단하다.

“타깃 모델의 기존 CP·PP 토폴로지는 그대로 둔 채, 시스템을 그 주위에서 적응시키면 드래프트 공동학습을 대규모·장문맥에서도 실용화할 수 있다.”

즉, 병렬화 구성을 드래프트에 맞춰 바꾸는 기존 접근과 달리, 타깃의 레이아웃을 보존하면서 두 메커니즘을 추가로 얹는다는 점이 이 논문의 독창성이다(근거: §2).

배경: 그들이 해결한 문제

RL 사후학습(GRPO 등)은 DeepSeekMath, DAPO, GLM5 등으로 대표되는 표준 패러다임이 됐고, 그 벽-클록 시간은 대부분 롤아웃 생성 이 차지한다(근거: §1). SD는 이를 병렬 검증으로 완화하지만, 기존 연구는 주로 드래프트 소싱·서빙 설정·검증 최적화에 집중했다(근거: §2, Related Work). 정작 “진화하는 정책을 따라 드래프트를 계속 학습시키는” 시스템 차원의 난제는 두 가지다.

1. Branch attention을 CP가 지원하지 않는다. EAGLE-3는 TTT(Train-Time Test)로 각 드래프트 위치마다 이전 예측 hidden을 재입력해 별도 branch 를 만들고, DFlash/DFlash 계열(DSpark)도 블록 내 위치들이 branch-local 컨텍스트에 어텐션해야 한다(근거: §2.2). 표준 causal CP는 이런 branch 구조를 표현하지 못한다. SpecForge가 EAGLE-3 TTT를 다뤘지만 순차 ring sharding으로 causal 워크로드가 불균형하고, Ulysses 차원에 제약이 걸린다(근거: §2.2).

2. PP에서 target feature가 드래프트와 멀리 떨어진 스테이지에 있다. 드래프트는 정책의 마지막 파이프라인 스테이지에 상주하는 서브모듈로 학습되는데, 드래프트가 소비하는 target hidden state(소위 tap) 는 앞선 여러 스테이지에서 생산된다. 표준 파이프라인 통신은 인접 스테이지만 연결하므로 이 비인접 feature를 전달하지 못한다(근거: §2.3).

저자들은 이 두 장애를 해결하는 엔드-투-엔드 시스템을 NeMo-RL에 통합했다(근거: §1).

새로운 접근법: Branch Attention under CP + TapChannel under PP

1) CP에서의 branch attention 분해

각 branch query는 두 개의 키 집합 에 어텐션해야 한다: (a) 주시퀀스의 causal prefix(CP 랭크들에 걸쳐 샤딩됨)와 (b) branch-local 키(branch의 anchor를 소유한 랭크에만 존재)이다. 저자들은 이 둘을 독립적으로 계산한 뒤 online-softmax(LSE) 감소로 병합한다(근거: §2.2, Fig. cp-branch-attention). 주시퀀스 컴포넌트는 표준 CP와 동일한 packed zigzag-ring attention 을 따르고, branch-local 컴포넌트는 anchor 소유 랭크에서 국소 계산된다.

여기서 $(O_{\mathrm{m}}, \ell_{\mathrm{m}})$ 는 주시퀀스, $(O_{\mathrm{b}}, \ell_{\mathrm{b}})$ 는 branch-local 성분의 출력과 log-sum-exp이다(근거: §2.2, Eq. 3). 이 하나의 병합 규칙이 EAGLE-3(TTT), DFlash(블록 병렬), DSpark(Markov head)를 동시에 커버한다는 점이 이 설계의 강점이다(근거: §2.2).

2) PP에서의 TapChannel

드래프트는 마지막 스테이지에만 있지만 tap을 생산하는 레이어는 여러 스테이지에 걸친다. tap은 회신 경로가 필요 없으므로, 저자들은 파이프라인 스케줄 밖의 사이드 패스 로 feature를 전달하는 TapChannel 을 설계했다(근거: §2.3, Fig. tap-channel). 드래프트 스테이지에 source별 메일박스(사전 할당 버퍼 슬롯)를 두고, source가 해당 마이크로배치의 policy forward를 마치면 tap을 슬롯에 쓰고, 드래프트는 같은 마이크로배치의 forward 직전에 읽는다. 순서 동기화는 슬롯마다 sequence stamp를 증분하는 방식으로 처리하며, 동일 노드 source는 CUDA IPC, 교차 노드 source는 GPUDirect RDMA가 붙은 전용 NCCL communicator를 쓴다(근거: §2.3).

작동 원리: 구체적인 예시로 살펴보기

CP branch attention 을 3×3 수준으로 축약해 보자. 시퀀스가 [A B C | D E F | G H I] 처럼 CP=3 랭크에 나뉘어 있다고 하자(근거: §2.2). EAGLE-3 TTT에서는 위치 C 뒤에, 드래프트가 스스로 예측한 hidden으로부터 branch 토큰 c1, c2 가 생성된다고 가정하자(근거: §2.2).

• 주시퀀스 성분: 각 랭크의 query는 자기보다 앞선 모든 주시퀀스 토큰을 봐야 한다. zigzag-ring에서 K/V가 랭크 사이를 순환하고 query는 제자리에 있으므로, 랭크가 D E F 를 맡고 있다면 A B C 의 K/V를 받아 어텐션한다. 이때 나온 출력을 $(O_{\mathrm{m}}, \ell_{\mathrm{m}})$ 라 하자.
• branch-local 성분: branch 토큰 c1, c2 의 anchor는 랭크 0(소유 A B C)이므로, c1, c2 의 키·밸류는 랭크 0에만 존재한다. 랭크 0만 국소적으로 branch 어텐션을 계산해 $(O_{\mathrm{b}}, \ell_{\mathrm{b}})$ 를 얻는다.
• 병합: 두 결과를 위 식으로 합친다. $\ell_{\mathrm{m}}=3.1$, $\ell_{\mathrm{b}}=1.2$ 라면 $\ell=\log(e^{3.1}+e^{1.2})\approx3.13$ 이 되고, 지수적으로 큰 성분(주시퀀스)이 최종 출력을 지배한다. 이는 ring attention 스텝 간 쓰는 online-softmax 감소와 정확히 같은 형태라 기존 커널 위에 자연스럽게 얹힌다(근거: §2.2).

통신 비용 관점 에서, CP 차수가 $C$ 이고 주시퀀스 토큰 $N$ 개가 랭크에 균등 분할되면 전방 ring에서 랭크당 송출 볼륨은

로, $d_{\mathrm{kv}}$ 는 토큰당 K/V 폭, $b$ 는 원소당 바이트다. branch-local K/V는 anchor 소유 랭크에 머물므로 이 비용은 branch의 수나 깊이와 무관 하다(근거: §2.2, Eq. 4). 어텐션이 문맥 길이에 2차로, 통신이 선형으로 증가하기 때문에 문맥이 길수록 통신을 계산 뒤로 숨길 여유(overlap headroom)가 커진다(근거: §2.2).

TapChannel 의 비용도 직관적이다. 마이크로배치의 토큰 행 수를 $n$, tap을 보내는 source 스테이지 집합을 $\mathcal{S}$, source $s$ 의 토큰당 feature 차원을 $d_s$ 라 하면 마이크로배치당 페이로드는

이며, 각 feature는 PP 홉 수와 무관하게 정확히 한 번만 직접 전송된다(근거: §2.3, Eq. 6). 파이프라인 스케줄이 source의 tap 생산과 드래프트 forward 사이에 자연스러운 여유(slack)를 주므로, 전송이 이 여유 안에 끝나면 지연이 0이다. 여유를 초과한 부분만 보이는 오버헤드가 된다(근거: §2.3, Eq. 7).

성능 검증: 주요 결과

1) 학습 궤적은 그대로 보존된다

가장 먼저 검증한 것은 “SD·공동학습이 RL 학습을 망치지 않는가"다. Qwen3-8B를 타깃으로 baseline(무 SD·무 공동학습)과 EAGLE-3/DFlash/DSpark 공동학습을 비교한 결과, reward·검증 정확도·training-inference KL 발산 모두 baseline 궤적을 밀착 추적했다(근거: §3.2, Fig. learning-stability). 특히 KL 발산이 0에 가깝다는 것은 훈련 백엔드와 추론 백엔드가 동일 가중치에서 수치적으로 일치함을 뜻하며, 이는 GRPO의 on-policy 가정을 지키는 핵심 조건이다(근거: §3.1).

2) 세 드래프트 계열 × 8B~122B 규모

공동학습된 드래프트는 수용 길이 2.28–4.78, 롤아웃 가속 1.19–2.23×, end-to-end 학습 가속 1.16–1.88× 를 달성했다(근거: §3.3, Tab. 1).

두 가지 패턴이 눈에 띈다(근거: §3.3). 첫째, DFlash/DSpark가 EAGLE-3보다 수용 길이가 길어 일관되게 빠르다. 둘째, 대형 MoE 타깃(Qwen3.5-122B, GPT-OSS-120B)은 수용 길이는 높지만 end-to-end 가속은 오히려 낮다 — 검증 forward마다 sparse routing이 더 많은 expert 연산을 유발하기 때문이다(근거: §3.3).

3) 멀티턴 워크로드의 한계

NeMo Gym Workplace Assistant(멀티턴·에이전틱·도구 사용 환경)에서 end-to-end 가속은 1.25–1.43× 로, 롤아웃 단계 가속 1.75–2.23×보다 크게 낮다. 그 이유는 롤아웃이 step 시간의 55.8% 밖에 안 되기 때문이다 — 나머지 도구 실행·환경 지연은 디코딩 가속으로는 건드릴 수 없다(근거: §3.3, Fig. workplace-learning).

4) CP 어텐션 성능: USP 대비 우위와 256K 확장

동일 워크로드에서 packed zigzag attention 을 SpecForge의 USP와 비교했다. USP는 배치를 실제 토큰 수의 2.25× 로 패딩하는 반면 packed 구현은 이 오버헤드를 피한다. CP=2, 4, 8에서 packed zigzag가 최선 USP보다 각각 2.9×, 2.3×, 1.5× 낮은 지연, 2.7× 낮은 GPU당 피크 메모리를 보였다(근거: §3.4, Fig. pack-dist). 장문맥 확장에서는 TTT 어텐션 지연이 CP=1의 17.7초에서 CP=8의 2.35초로 7.5×(병렬 효율 94%) 줄었고, GPU당 메모리는 53.2 GB에서 7.5 GB로 거의 선형 감소했다(근거: §3.4, Fig. cp-scaling).

5) PP 오버헤드는 “겹쳐지는” 수준

TapChannel의 원시 전송은 one-sided write로 27–39 GB/s, host staging 대비 4.5–8.5× 빠르며, source 스테이지는 노이즈 수준, 수신 드래프트 스테이지는 1.6% 의 HBM 경합만 겪는다. 반면 host staging은 모든 랭크를 80% 이상 느리게 한다(근거: §3.5, Fig. tap-transport). 전체 실행에서 드래프트 공동학습은 block 드래프트 기준 15% 미만, EAGLE-3는 34%(TTT 패스 탓)의 업데이트 시간 오버헤드를 더한다. 그러나 롤아웃을 28–52% 단축해 net speedup 1.31–1.85× 를 냈고, Tap 대기 시간은 정책 업데이트당 0.4–0.6초(최적화 시간의 1.5–2.2%)에 불과해 파이프라인 스케줄에 거의 숨겨진다(근거: §3.5, Tab. 2).

우리의 관점: 강점, 한계, 그리고 이 연구가 중요한 이유

강점 은 명확하다. 첫째, 타깃 토폴로지 불변이라는 설계 원칙이 돋보인다. 기존 연구가 병렬 레이아웃을 드래프트에 맞춰 바꾼 반면, 이 논문은 정책의 CP·PP 구성을 건드리지 않고 시스템을 적응시켜 “정책을 계속 학습시키는 와중에” 드래프트도 함께 학습할 수 있게 했다(근거: §2, Related Work). 둘째, 하나의 LSE 병합 규칙이 EAGLE-3·DFlash·DSpark 세 가지 상이한 드래프트 구조를 동시에 지원하는 일반성 이다(근거: §2.2). 셋째, 8B→122B, 단일턴→멀티턴에 걸친 폭넓은 검증과 함께 학습 궤적 보존을 KL 발산으로 정량화한 점이 실무 신뢰를 높인다(근거: §3.2).

한계 도 분명하다. 첫째, MoE 타깃과 linear-attention 모델에서는 speculative decoding의 이득 자체가 구조적으로 제한된다 — MoE는 검증마다 expert 라우팅 비용이, linear attention은 검증 자체가 이미 저렴해 상대적 가속 여지가 작다(근거: §3.3, §4). 둘째, EAGLE-3의 TTT 다중 패스가 34%의 업데이트 오버헤드를 초래해, block 계열 드래프트에 비해 불리하다(근거: §3.5, Tab. 2). 셋째, 멀티턴·도구 사용 환경처럼 롤아웃이 전체 step의 절반 정도만 차지하면, 디코딩 가속의 end-to-end 전이율이 크게 떨어진다는 점은 이 접근의 가치가 롤아웃 지배적 워크로드에 집중 됨을 시사한다(근거: §3.3). 또한 이 논문은 “시스템을 만든다"는 데 무게가 실려 있어, 드래프트 아키텍처 자체의 새 기여보다는 기존 드래프트를 대규모 분산 학습에 실용화하는 기여가 본질이다.

그럼에도 이 연구가 중요한 이유는, RL 사후학습의 실제 병목이 “모델 학습"이 아니라 “롤아웃 생성"이라는 점을 시스템 레벨에서 정공법으로 파고들었기 때문이다. 학습 궤적을 보존하면서 드래프트를 함께 학습시키는 체계가 없다면, SD 기반 RL 가속은 소규모·단문맥에서만 가능한 일회성 트릭에 머물렀을 것이다.

다음 단계는?: 앞으로의 길

저자들은 sparse MoE와 linear-attention 모델에 맞게 speculative decoding을 조정하는 것을 향후 과제로 꼽는다(근거: §4). 이에 더해 합리적인 다음 단계를 제안해 보면:

• MoE 검증 비용 절감: 검증 forward의 expert 연산을 줄이는 라우팅 캐싱이나, expert 수준의 speculative 검증을 결합해 대형 MoE에서도 수용 길이 이득을 end-to-end 가속으로 전환.
• 멀티턴 롤아웃의 미디코딩 병목 완화: 도구·환경 지연 구간을 디코딩과 겹치거나, 멀티턴 구조에서 acceptance를 높이는 세션 단위 드래프트 컨디셔닝.
• EAGLE-3 TTT 오버헤드 축소: TTT 패스 수를 줄이거나 DSpark식 저비용 Markov head를 접목해 block 계열 수준의 업데이트 비용으로 수렴시키기.
• CP 강한 확장에서의 통신 노출 완화: 문맥이 짧아질 때 통신이 병목이 되는 지점(§2.2)을 개선하는, branch-aware한 K/V 압축 또는 셔플 일정 최적화.

결국 이 논문의 의의는 “특정 드래프트 하나"가 아니라, 진화하는 정책과 함께 살아 움직이는 드래프트를 대규모 분산 학습에서 실용화하는 패턴 을 제시했다는 데 있다.

논문 원문의 표

arXiv e-print 의 LaTeX 원본에서 기계적으로 옮긴 표입니다. 숫자는 논문의 값이며 모델을 거치지 않았습니다.

표 1. End-to-end performance across target scales and draft families.

표 2. Pipeline-parallel overhead for Qwen3-8B, averaged over first 10 policy update steps. DSpark achieves the largest end-to-end speedup (1.85$\times$) with 14.6% update-time overhead.

이 글의 그림은 arXiv:2609.07108 원본에서 가져왔습니다 (CC BY 4.0). 크기와 형식만 바꿨습니다.