Online co-training of draft models to speed up speculative decoding in large-scale RL post-training

Post-training with reinforcement learning (RL) has become a decisive phase in pushing large language models towards more sophisticated reasoning behaviours, but it has a well-known bottleneck: rollout generation, that is, the text sequences the model produces during training that are then evaluated to compute the reward signal. This paper, a technical report published on arXiv by Zili Wang, Zhaopeng Qiu, Yuekai Zhang, Shuang Yu and Junjie Lai, tackles precisely that bottleneck through speculative decoding, a technique already established in the inference world that is taken here into the more demanding territory of large-scale, long-context RL training.

Speculative decoding works by relying on a much smaller and faster 'draft' model that proposes several candidate tokens, which the large model (the policy being trained) verifies in parallel instead of generating them one by one autoregressively. The more proposed tokens the large model accepts, the greater the speedup obtained. The paper notes that online co-training of the draft —that is, continuously updating that small model as the policy evolves during RL— improves its predictive accuracy and therefore increases those speedups even further, rather than using a static draft that progressively falls out of sync with the model that is learning.

Scaling this co-training to large models with long contexts, however, presents two concrete technical obstacles that the authors identify explicitly. The first concerns context parallelism (CP), the technique that allows very long sequences to be split across several GPUs. Standard implementations of causal attention under context parallelism (such as zigzag ring attention) do not support so-called 'branch attention', a tree-shaped attention pattern needed to verify several branches of speculatively generated candidate tokens simultaneously. The second obstacle arises in pipeline parallelism (PP): the 'target features' —the large model's intermediate representations that the draft needs in order to train and stay aligned with it— end up spread across different pipeline stages, which complicates transporting them without disrupting the rest of the training process.

To solve the first problem, the team extends the packed, load-balanced zigzag ring attention scheme, fusing per-GPU-rank local branch attention with the causal attention of the main sequence, so that both types of attention coexist within the same context-parallelism scheme without breaking its efficiency. For the second problem, they propose a mechanism called TapChannel, which transports those target-model intermediate features between pipeline stages over a separate communication path, leaving the usual pipeline schedule intact and thus avoiding interference with the rest of training.

According to the experiments described in the abstract, draft models co-trained this way manage to track the base policy closely —that is, they maintain good predictive fidelity despite being constantly updated— while delivering substantial speedups both in the rollout generation phase and in end-to-end training, at model scales of up to 122 billion parameters (122B). The proposed context-parallelism design also shows solid scaling at contexts of up to 256,000 tokens, with significant memory savings over prior work, while the pipeline-parallelism transport mechanism (TapChannel) introduces only modest overhead.

The work falls under arXiv's machine learning (cs.LG) and distributed, parallel and cluster computing (cs.DC) categories, was submitted on 7 September 2026, and the authors state that the code is publicly available, although the downloaded content does not specify the repository link or any implementation details beyond what is summarized in the abstract.

Beyond the specific technical details, the paper is representative of a broader trend in language model training infrastructure: as RL post-training —especially that aimed at extended reasoning and very long contexts— becomes central to frontier model development, optimizing the computational cost of the rollout generation phase stops being a concern exclusive to production inference and becomes a critical factor within the training loop itself. Integrating acceleration techniques such as speculative decoding directly into distributed RL pipelines, resolving the specific frictions introduced by context parallelism and pipeline parallelism at scale, is the kind of systems engineering work that, while less visible than the launch of a new model, directly determines how much it costs and how long it takes to train the models that do end up making the news.