PipelineRL: Optimizing LLM Reinforcement Learning via Inflight Weight Updates

PipelineRL is an experimental reinforcement learning (RL) implementation designed to solve the fundamental trade-off between high inference throughput and on-policy data collection in large-scale LLM training. By implementing inflight weight updates, PipelineRL allows inference servers to receive updated model weights without stopping the inference process, ensuring high GPU utilization while keeping training data near on-policy.

Solving the Inference Throughput vs. On-Policy Trade-off

In conventional RL workflows, there is a conflict between efficiency and data freshness. To achieve high throughput, inference servers typically use large batch sizes, which generate data for multiple policy optimization steps. However, each subsequent optimization step increases the "lag" between the weights used to collect the data and the current policy weights, making the data increasingly off-policy and less effective for training.

PipelineRL remediates this by updating weights in inference servers after each optimizer step without stopping inference. The system only pauses inference servers for the brief duration required to receive new weights. This approach allows the inference server to maintain optimal batch sizes while ensuring data remains on-policy or near on-policy, leading to more stable and effective learning.

Performance and Stability Results

Experiments using 7B and 32B models trained on the Open-Reasoner-Zero dataset demonstrate that PipelineRL matches or exceeds the performance of Open-Reasoner on the AIME 2024 and MATH 500 reasoning benchmarks.

Simplified RL Algorithm

Despite its competitive performance, PipelineRL uses a significantly simpler RL implementation than Open-Reasoner-Zero. Key simplifications include:

  • Simplified GRPO: It uses a simplified version of Group Relative Policy Optimization (GRPO) without a value function.
  • No Complex Filtering: The implementation omits trust region importance weight clamping, overlong sequence filtering, and reward shaping.
  • Basic Loss Normalization: Loss is normalized using the number of sequences in the batch as the denominator, giving equal weight to all tokens.
  • No Penalties: The system uses no KL penalty or entropy bonus (though reference model KL is supported).

Impact of KV Cache Stale Data

A primary concern with inflight weight updates is that sequence generation continues with stale keys and values in the KV cache, as these were computed with a previous version of the model. However, experimental results indicate that this does not adversely affect training stability.

Modular Architecture and Technical Contracts

PipelineRL is designed to be modular to allow integration with specialized inference (e.g., SGLang, vLLM) and training (e.g., DeepSpeed, FSDP, TorchTitan) software. This is achieved through two primary contracts:

Inference Contract

To integrate with PipelineRL, inference software must expose three specific APIs:

  1. Process group initialization: An HTTP POST /init_process_group request to initialize the process group for weight updates.
  2. Weight Update Trigger: An HTTP POST /request_weight_update request that signals inference servers to pause and receive a weight broadcast via NCCL.
  3. Chat completion: Standard HTTP POST /v1/chat/completion requests for actor interaction.

Trainer Contract

Training software must provide Python APIs for the following operations:

  • Worker initialization: Loading and sharding training weights and optimizer state.
  • Forward pass: Producing token log-likelihoods.
  • Backward step: Computing and accumulating gradients for the RL objective.
  • Optimizer Step: Executing the optimizer step.
  • Weight gathering and broadcasting: Gathering updated weights layer-by-layer for broadcasting to inference servers.

Experimental Configuration

PipelineRL was tested on 7B and 32B models with the following hyperparameters:

  • Batch size: 4096
  • Learning rate: 1e-6
  • Max generated tokens: 8192

Training compute requirements were approximately 3.5 days on 2 nodes for the 7B model and 6 days on 4 nodes for the 32B model.

Sources