TL;DR

NVIDIA vừa công bố cách tích hợp speculative decoding vào pipeline RL post-training mà không làm thay đổi behavior của model. Kết quả: 1.8× tốc độ sinh token ở 8B, chiếu 2.5× tăng tốc end-to-end ở 235B scale. Điều làm cho phương pháp này khác biệt là tính lossless - policy vẫn học từ chính xác distribution của model gốc.

NVIDIA Accelerating RL Post-Training via Speculative Decoding

Tại sao sinh token lại là nút thắt cổ chai

Trong RL post-training (RLHF, GRPO, PPO...), vòng lặp training gồm hai giai đoạn: rollout (model tự sinh câu trả lời) rồi cập nhật trọng số theo reward. Vấn đề: rollout chiếm 65-72% thời gian mỗi step trong synchronous RL pipeline.

Tốc độ sinh token phụ thuộc vào latency per token - autoregressive generation sinh từng token một, mỗi token phải chờ kết quả của token trước. Đây là điểm mà speculative decoding có thể can thiệp.

Speculative decoding hoạt động như thế nào

Ý tưởng cốt lõi: dùng một draft model nhỏ hơn để dự đoán nhiều token tiếp theo cùng lúc, rồi target model lớn xác nhận song song. Nếu dự đoán đúng, hệ thống chấp nhận nhiều token trong một bước verify.

Điều quan trọng với RL: quá trình rejection sampling của speculative decoding được chứng minh toán học là lossless - sequence cuối cùng có cùng distribution như thể target model tự sinh ra. Không có distribution mismatch, không cần off-policy correction, không làm thay đổi training signal.

Đây là lợi thế quyết định so với các phương pháp tăng tốc khác:

  • Async execution: policy lag gây distribution shift, cần correction

  • Off-policy replay: đánh đổi fidelity lấy throughput

  • Low-precision rollouts: thay đổi output distribution

  • Speculative decoding: tốc độ cao + distribution nguyên vẹn

Con số thực tế

Metric

RL-Zero (8B)

RL-Think (8B)

235B (projected)

Generation speedup

1.8×

1.5×

3.0-3.5×

Overall step speedup

1.41×

1.35×

~2.5×

Generation time before

100s

133.6s

-

Generation time after

56.6s

87.0s

-

EAGLE-3 acceptance length

3.32 tokens

2.77 tokens

-

Hardware: 32 GPU GB200 (8-node, 4 GPU/node). Projection 235B chạy trên 2048 GPU GB200.

Lưu ý quan trọng: tại sao overall step speedup (1.35-1.41×) thấp hơn generation speedup (1.5-1.8×)? Vì log-probability recomputation và training phase chưa được accelerate - chúng chiếm 28-35% còn lại của step time.

EAGLE-3 và những điều bất ngờ

NVIDIA dùng EAGLE-3 làm draft model framework - model-agnostic, hoạt động với bất kỳ pretrained model nào. Một số phát hiện đáng chú ý:

  • k=3 là tối ưu, không phải k=5 hay k=7. Draft dài hơn tạo overhead verification vượt qua lợi ích - RL-Think thậm chí chậm hơn baseline khi k≥5.

  • N-gram drafting thất bại: dù acceptance length dương (2.05-2.47 tokens), n-gram khiến hệ thống chậm hơn autoregressive (0.5-0.7×). Acceptance length ≠ speedup.

  • Domain-specific initialization quan trọng: draft model init từ domain phù hợp cho 1.77× vs 1.51× từ generic chat data.

  • Online adaptation: giúp draft model yếu theo kịp policy đang thay đổi; ít tác dụng khi init đã tốt.

Lossless - tại sao đây mới là điểm then chốt

Trong RL, model học bằng cách sinh ra hành vi (rollout), nhận reward, rồi update policy. Nếu rollout bị tạo ra từ distribution khác với policy hiện tại - dù chỉ một chút - gradient update sẽ không còn chính xác. Đây là lý do mà các phương pháp tăng tốc khác đều phải đánh đổi:

  • Async RL cho phép policy lag: rollout cũ được dùng để update policy mới - cần importance sampling correction

  • Off-policy replay: recycle experience cũ - nặng về correction overhead

  • Low-precision rollouts: thay đổi token probabilities - ảnh hưởng tới KL penalty và logprob

Speculative decoding không có vấn đề này. Rejection sampling procedure được chứng minh toán học: dù draft model đề xuất token nào, xác suất chấp nhận được tính sao cho output cuối cùng có đúng distribution của target model. Policy học từ output của chính mình - không cần correction nào thêm.

Ai nên chú ý

Phương pháp này đặc biệt có giá trị cho:

  • Team đang chạy RL post-training cho reasoning model (math, code, verifiable tasks) - đây là workload mà generation chiếm đa số thời gian

  • Frontier model lab vận hành cluster 2048+ GPU ở scale 235B+ - 2.5× speedup nghĩa là cắt đáng kể compute cost

  • Bất kỳ team nào dùng NeMo RL + vLLM stack - tích hợp có sẵn trong v0.6.0

Hạn chế cần biết: hiện tại chỉ chạy trong NeMo RL; portability sang các framework khác (Axolotl, TRL...) chưa được xác nhận. Draft model cần được giữ đồng bộ với policy đang evolve.

Tiếp theo

Paper không công bố roadmap cụ thể, nhưng có hai hướng nghiên cứu tự nhiên: (1) accelerate log-prob recomputation để đẩy overall speedup gần hơn với generation speedup, và (2) port integration sang các RL framework phổ biến hơn như TRL.

Code open source, Apache 2.0, có sẵn tại github.com/NVIDIA-NeMo/RL. Paper: arXiv 2604.26779.

Via: arXiv paper, The Turing Post.