PipelineRL

We are excited to open-source PipelineRL, an experimental RL implementation that tackles a fundamental challenge in large-scale Reinforcement Learning with LLMs: the trade-off between inference throughput and on-policy data collection. PipelineRL's key innovation is inflight weight updates during RL training (see Figure 1 below). This allows PipelineRL to achieve constantly high inference throughput and minimize the lag between the weights used for rollouts and the most recently updated model weights. The result: fast and stable RL training for large language models.

In this blog post, we show that 1) inflight weight updates do not harm the training process and 2) PipelineRL achieves competitive results compared to Open-Reasoner-Zero, while using a simpler RL algorithm. We also present the modular PipelineRL architecture that facilitates trying new inference / trainer combinations.

In conventional RL approaches (Figure 1a), there is a trade-off between high throughput inference and on-policy data collection. To explain this trade-off let us first define conventional RL algorithmically:

PipelineRL (Figure 1b) remediates this trade-off through inflight weight updates. We update the weights in inference servers after each optimizer step without ever stopping inference . We only pause the inference at all inference servers for just the time needed to receive the new weights. Inflight weight updates allow the inference server to constantly maintain the optimal batch size while simultaneously ensuring data remains on-policy or near on-policy, which leads to better GPU utilization and more effective learning, respectively.

Notably, our RL implementation is much simpler than Open-Reasoner-Zero. While Open-Reasoner-Zero uses a value function, our implementation is a simplified version of GRPO. In particular, we found that the trust region importance weight clamping is not needed for stable training. Neither was overlong sequence filtering or reward shaping from the DAPO paper . For normalizing the loss we just use the number of sequences in the batch as the denominator, giving equal weights to all tokens. We used no KL penalty and no entropy bonus (though our implementation does support reference model KL). Despite the simplicity of our implementation, or perhaps thanks to it, training is very stable as you can see in this wandb report .

More research coming soon. More analysis is needed to understand how inflight weight updates affect the training dynamics, and to carefully measure the speed-ups that PipelineRL brings. Also, much can be said about the similarities between PipelineRL and highly relevant prior work on asynchronous Reinforcement Learning for LLMs. For all this and more please stay tuned for our upcoming research paper!

Upcoming features. Our implementation is still experimental and lacks some important functionality. Top priorities for us include using coroutines for more precise inference batch size control, multi-modal support and sequence parallel training. We would also welcome contributions of more inference server and trainer integrations. We will not, however, try to make the pipeline-rl repo a framework that supports all possible algorithms and reward functions. Our take is that pipeline-rl should be a hackable and fast reference implementation of GRPO with easily verifiable rewards. If you'd like to do a research project using PipelineRL, you can just fork the repo and have fun hacking the code!

To achieve high throughput, the inference servers must use large batch sizes and, therefore, generate data for multiple policy optimization steps. However, each optimization step increases the lag between the current policy and the data collected using the inference policy, progressively making collected data more off-policy and less effective for training. On-policy learning requires data for a single optimization step. But producing small amounts of data with many GPUs is inefficient because this means the per-GPU batch size is small. Furthermore the batch size goes down as the inference server finishes the short sequences and has only the few longest sequences in progress.

To demonstrate the effectiveness of PipelineRL and the benefits of inflight weight updates, we trained a 7B model and 32B model on the Open-Reasoner-Zero dataset. Looking at the learning curves we see that PipelineRL matches or exceeds the performance of Open-Reasoner on the popular reasoning test benchmarks: AIME 2024 and MATH 500 (see Figure 2 above).

One might expect that inflight weight updates would result in an unstable training process, since sequence generation proceeds with stale keys and values in the KV cache that was computed with a previous model version. However, our experiments indicate this does not adversely affect stability.