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

Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training. By Zili Wang and four others, from NVIDIA. Published on the seventh of September. Its code repository is worked out: about one thousand five hundred source files, with a readme and a licence.

In reinforcement-learning post-training, generating rollouts dominates the cost. Speculative decoding speeds that generation up, and co-training the draft model online, alongside the policy, makes the draft more accurate and the speedups larger. Accuracy in the draft is what converts into speed. Scaling online co-training to large models with long contexts runs into two obstacles. Branch attention is not supported by standard causal context- parallel, or CP, implementations. And the target features the draft needs are spread across pipeline-parallel, or PP, stages. Both are the parallelism schemes that large, long-context training relies on, so the co-training has to work within them rather than around them.

The authors address both with an end-to-end system for large-scale online draft co-training. Each obstacle gets its own fix, one for context parallelism and one for pipeline parallelism.

For context parallelism, they extend packed, load-balanced zigzag ring attention by merging rank-local branch attention with the causal attention over the main sequence. For pipeline parallelism, a component called TapChannel carries intermediate target features across stages along a separate path, so the pipeline schedule itself is left untouched. The features reach the draft without the pipeline having to be rearranged to deliver them.

Co-trained drafts closely track the policy baseline, and they deliver substantial speedups, both in rollouts and end to end, across model scales up to one hundred and twenty-two billion parameters. The CP design scales strongly at two hundred and fifty-six thousand tokens, with significant memory savings over prior work. The PP transport adds only modest overhead. Tracking the policy baseline closely is the condition that matters: speed gained from the draft should not come at the cost of what the training produces.

Neither fix changes what co-training is. They make it fit inside the parallelism that large, long-context training already depends on: attention merged into the existing ring scheme, and features moved beside the pipeline rather than through it. Code is linked from the NVIDIA-NeMo repository.