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.

Conventional RL vs PipelineRL

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:

current_policy = initial_policy
opt_state = init_optimizer(current_policy)

while True:
    # RL step starts
    # inference
    inference_policy = current_policy
    list_of_prompts = [sample_prompts(training_batch_size) \
        for _ in range(num_grad_steps)]
    list_of_rollouts = [sample_rollouts(prompts, inference_policy) \
        for prompts in list_of_prompts]
    # training
    lag = 0 # lag between the inference and current policies
    for rollouts in list_of_rollouts:
        current_policy, opt_state = policy_update(current_policy, opt_state, rollouts)
        lag += 1
    # RL step ends