Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training
NVIDIA researchers detail an end-to-end system for online draft co-training that speeds speculative decoding in large-scale long-context RL post-training.
The paper tackles scaling online draft co-training for speculative decoding in RL post-training, where rollout generation dominates cost. It extends packed, load-balanced zigzag ring attention to merge rank-local branch attention with causal main-sequence attention for context parallelism, and introduces TapChannel to transport target features across pipeline-parallel stages without changing the schedule. Experiments show co-trained drafts tracking the policy baseline with substantial rollout and end-to-end speedups up to 122B parameters and strong scaling at 256K tokens.
- Extends packed zigzag ring attention with rank-local branch attention for causal context parallelism.
- TapChannel transports target features across pipeline-parallel stages without altering the schedule.
- Co-trained drafts track the policy baseline with substantial rollout and end-to-end speedups.
- Validated up to 122B parameters and 256K context with memory savings.
- Tags
- context-parallelisminference-optimizationnvidiapipeline-parallelismreinforcement-learningspeculative-decoding
- Vendors
- NVIDIA
- Products
- NVIDIA NeMo RL
- Organizations
- NVIDIA
From the original article
1 min read at huggingface.co
Speculative decoding accelerates rollout generation, which dominates the cost of reinforcement learning (RL) post-training. Online co-training can further increase the draft's accuracy, yielding greater speedups.
By Zili Wang, Zhaopeng Qiu, Yuekai Zhang · huggingface.co
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.