PipelineRL: Faster On-policy Reinforcement Learning for Long Sequence Generation

Thank you for your insightful review. We first highlight our new experiments and then addresses your comments one by one.

We provide new experiments to addresses 1) scalability, and 2) data quality and stability.

First, we improve the stability fo the code base by adding a value function. The value function resulted in Conventional RL being stable up to G=16 from G=4, making it a much stronger baseline. Furthermore, we scaled PipelineRL to 128 GPUs and still observe a significant throughput improvement over the baseline.

Second, in order to understand the impact of inflight weight update, we compute the KL with the target policy ($\pi$) for different max lag for the conventional RL behavior policy ($\pi_C$) and PipelineRL behavior policy ($\mu$) for 3 different checkpoints. Interestingly, we oberve that PipelineRL generally obtains lower KL divergence between the behavior and target policy for max lag greater than 1.

The continuous "in-flight weight updates" approach could lead to training instabilities, especially since behavior policies partly depend on outdated cached activations (KV vectors). Although the authors propose an ESS-based monitoring safeguard, it's unclear how comprehensively this mitigates instability under rapid or extensive model updates.

In order to assess the data quality produced by PipelineRL, we added an experiment (see Exp 2 above) to compute the KL divergence between the policy with inflight weight update and the target policy. We observe that the inflight weight update policy obtains similar or lower KL divergence than Conventional RL for the same maximum lag. This implies that weight updates can improve the on-policiness of the collected data and its quality.

While scalability is demonstrated convincingly up to 32 GPUs, the paper does not discuss potential scalability limitations at larger scales (e.g., hundreds of GPUs). Issues such as communication bottlenecks, KV-cache complexity, or diminishing returns in throughput at very large scales are plausible and warrant further investigation.

We successfully scaled PipelineRL to 128 GPUs and show significant speed improvement (see Exp 1 above). As mentionned in Figure 3, the major concern in scaling up the number of GPUs is the max lag will increase and can potentially destabilize training.

PipelineRL's performance improvements hinge on carefully tuned hyperparameters (e.g., ratio of training to generation GPUs, lag, batch sizes). The paper lacks a comprehensive sensitivity analysis, making it difficult to assess how robust the method is to variations in these critical parameters.

In our experience, allocating 50% of your GPUs for training and 50% for inference is a good starting point. One can then monitor if the trainer is starved or overrun with data and adjust the allocation accordingly.

Since PipelineRL specifically targets long-sequence generation, it would be valuable to see analysis or discussion on how variations in sequence lengths affect PipelineRL's throughput, effectiveness, and stability.

• Please see appendix A3, for a theoretical derivation of how the max sequence length impact the lag and constain throughput and appendix A4 for a theoretical speed up computation.

Questions:

Could the authors provide further insights into specific instability scenarios arising from continuous asynchronous updates? Are there recommended heuristics or additional safeguards beyond ESS-based monitoring?

In order to assess the potential instability caused by using inflight weight updates, we added an experiment (see Exp 2 above) to compute the KL divergence between the policy with inflight weight update and the target policy. We observe that the inflight weight update policy obtains similar or lower KL divergence than Conventional RL for the same maximum lag. This implies that weight updates can improve the on-policiness of the collected data and its quality.

Have you explored or observed scalability limitations beyond the demonstrated 32-GPU setup? What are potential bottlenecks at larger scales?

Please see experiment 1 above. We successfully scaled PipelineRL to 128 GPUs.

Could the authors provide additional guidance or insights into selecting optimal PipelineRL configurations (e.g., number of accelerators, optimal batch sizes, lag settings)?

Please see Appendix A3 for a dicussion of PipelineRL as a function of accelerators, batch size, and the resulting lag.

Could the authors discuss how PipelineRL's performance and stability are influenced by variations in sequence length? Are there specific considerations or adjustments recommended when dealing with significantly shorter or longer sequences than those tested?

On the one hand, for very short sequences, e.g. 1 token, PipelineRL would not provide much improvement over Conventional RL (for high enough G) since the inference GPU would also run at full batch size. On the other hand, for very long sequences, e.g. 1M tokens, the performance gains over conventional RL will be large, but the max lag will increase and could potentially destabilize learning. However, lowering the learning rate, increasing the batch size, and excluding tokens with very high lag could potentially be used to stabilize training and still obtain good throughput.

Thank you for your review. We first highlight our new experiments and then addresses your comments one by one.

We provide new experiments to addresses 1) scalability, and 2) data quality and stability.

First, we improve the stability fo the code base by adding a value function. The value function resulted in Conventional RL being stable up to G=16 from G=4, making it a much stronger baseline. Furthermore, we scaled PipelineRL to 128 GPUs and still observe a significant throughput improvement over the baseline.

Second, in order to understand the impact of inflight weight update, we compute the KL with the target policy ($\pi$) for different max lag for the conventional RL behavior policy ($\pi_C$) and PipelineRL behavior policy ($\mu$) for 3 different checkpoints. Interestingly, we oberve that PipelineRL generally obtains lower KL divergence between the behavior and target policy for max lag greater than 1.

Testing on mathematical reasoning model alone seems to be slightly underwhelming for the potential impact of this work. I am wondering if it including other tasks e.g. code generation would help to see generalisation of PipelineRL to other tasks?

Training a model from scratch to improve mathematical reasoning is particularly difficult since the sequence lengths vary acros problems and throughout training. The improvement will translate to other domains also having variablility in sequence lengths.

It would be helpful to include results from multiple runs and show variations e.g. standard deviation, to support the current results with statistical significance.

Yes, we have ran 3 seeds for experiment 1 and are in the process of running 2 more seeds.

Do you anticipate similar performance improvements when applying PipelineRL to other domains like code generation and dialog?

PipelineRL improvement over conventional RL are partly a function of the variance in sequence length. Therefore, a code/dialog dataset that results in sequence of different length would also benefit from PipelineRL.

Is there any plans to perform ablations to isolate the contribution of in-flight updates vs. asynchronous pipeline parallelism to the overall speedup?

Removing in flight weight update would result in stopping the inference GPU and result in drastic slowdown of PipelineRL.

Thank you for your review.

it would be valuable to include measurements of potential degradation in data quality during the PipelineRL training process. This would provide additional insights into the trade-offs introduced by the proposed approach.

In order to understand the impact of inflight weight update, we compute the KL with the target policy ($\pi$) for different max lag for the conventional RL behavior policy ($\pi_C$) and PipelineRL behavior policy ($\mu$) for 3 different checkpoints. Interestingly, we oberve that PipelineRL generally obtains lower KL divergence between the behavior and target policy for max lag greater than 1.

We also improve the stability fo the code base by adding a value function. The value function resulted in Conventional RL being stable up to G=16 from G=4, making it a much stronger baseline. Furthermore, we scaled PipelineRL to 128 GPUs and still observe a significant throughput improvement over the baseline.

Thank you for your review. We first highlight our new experiments and then addresses your comments one by one.

We provide new experiments to addresses 1) scalability, and 2) data quality and stability

First, we improve the stability fo the code base by adding a value function. The value function resulted in Conventional RL being stable up to G=16 from G=4, making it a much stronger baseline. Furthermore, we scaled PipelineRL to 128 GPUs and still observe a significant throughput improvement over the baseline.

Second, in order to understand the impact of inflight weight update, we compute the KL with the target policy ($\pi$) for different max lag for the conventional RL behavior policy ($\pi_C$) and PipelineRL behavior policy ($\mu$) for 3 different checkpoints. Interestingly, we oberve that PipelineRL generally obtains lower KL divergence between the behavior and target policy for max lag greater than 1.

There are no major weaknesses, but significance could be improved by demonstrating the method on more models, in more environments, and against more baselines.

Training a model from scratch to improve mathematical reasoning is particularly difficult since the sequence lengths vary acros problems and throughout training. The improvement will translate to other domains also having variablility in sequence lengths and other autoregressive models.

Significance could also be improved by further clarifying the types of regimes in which the proposed approach is expected to be effective.

We provided extensive discussion of the throughput and lag of PipelineRL as a function of max sequence length, batch size, and number of accelerators in Appendix A3.

1. Just curious: In Figure 5b, does the red line shift closer to the blue if G=8? I.e., can the speed gap be largely addressed by increasing G, and the main concern is how off-policy you become with large G?

Yes, this is what we refer to as the trade-off between hardware efficiency and data on-policyness for LLM training. Increasing G improves hardware efficiency at the cost of more off-policy data.

• Line 126 and Figure 3a's caption: stating some assumptions could help clarify when this analysis holds

Figure 3 is a didactic example assuming that doubling the number of accelerators will double the throughput of the inference GPUs and of the training GPUs. Under these assumptions, for a sequence of length L, the lag of the earliest token will double.

• Add training time to Table 1

We will add the training time for our models.

3. Major clarity concern. Consider adding more clarity around when PipelineRL provides benefits (in terms of sequence lengths, GPU counts, etc.). For example, what are the speedups as a function of max sequence length?

• Please see appendix A3, for a theoretical derivation of how the max sequence length impact the lag and constrain throughput and appendix A4 for a theoretical speed up computation.

4. Baseline choice concern:
   • Compare PipelineRL's speed and performance to those of asynchronous RL approaches

In generation-bound RL scenarios like math reasoning, where sequence generation take significantly longer than training, async (on-policy) RL results in highly inefficient training GPUs utilization. The efficiency of async RL can be improved by using off-policy methods, i.e. training several time on the same data. In contrast, Conventional RL takes better advantage of the GPUs by using all the GPUs for either training or generating without the need to training several time on the same data.

• Demonstrate the effectiveness of PipelineRL with another model (beyond Qwen)

Qwen is currently the most popular model to train from scratch for math reasoning. Benchmarking PipelineRL on training Qwen from scatch makes it easy for the community to compare and evaluate our results.