Asynchronous RLHF: Faster and More Efficient Off-Policy RL for Language Models

Authors: Michael Noukhovitch, Shengyi Huang, Sophie Xhonneux, Arian Hosseini, Rishabh Agarwal, Aaron Courville
Primary area: foundation or frontier models, including LLMs
Keywords: reinforcement learning from human feedback, efficient llm finetuning, off-policy RL

Abstract:
The dominant paradigm for RLHF is *online* and *on-policy* RL: synchronously generating from the large language model (LLM) policy, labelling with a reward model, and learning using feedback on the LLM's own outputs. While performant, this paradigm is computationally inefficient. Inspired by classical deep RL literature, we propose separating generation and learning in RLHF. This enables asynchronous generation of new samples while simultaneously training  on old samples, leading to faster training and more compute-optimal scaling. However, asynchronous training relies on an underexplored regime, online but *off-policy* RLHF: learning on samples from previous iterations of our model which give a worse training signal. We tackle the fundamental challenge in this regime: how much off-policyness can we tolerate for asynchronous training to speed up learning but maintain performance? Among several RLHF algorithms we test, online DPO is found to be most robust to off-policy data, and robustness increases with the scale of the policy model. We study further compute optimizations for asynchronous RLHF but find that they come at a performance cost, giving rise to a trade-off. We verify the scalability of asynchronous RLHF by training a general-purpose chatbot from LLaMA 3.1 8B on an instruction-following task $\sim$40\% faster than a synchronous run while matching final performance. Finally, we extend our results to math and reasoning to demonstrate asynchronous RL can finetune Rho 1B on GSM8k $\sim$70\% faster while matching synchronous accuracy.

Decision: accept (poster)
Reviewer scores: 5, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729905192515

thread_type:
official_review

<thread_object>
creator: reviewer_rlxt

<content>
summary:
This paper studies incorporating off-policy asynchronous training to speed up RLHF. The main idea is to use separate GPUs for generation and training, which can leverage the acceleration techniques developed for both training and inference. Experiment results show that DPO is more robust under off-policy data, and the proposed method can achieve 40% speed up in training LLaMA 3.1 8B without sacrificing final performance.

strengths:
This paper studies an important problem: how to speed up RLHF training. The proposed method is easy to implement and the experiment results are quite promising. I think it is quite meaningful to study such off-policy asynchronous training setting, as fully online training is quite expensive and leads to low utilization of GPUs. This paper gives a thorough study of such setting with meaning conclusions.

weaknesses:
1. There are many other RLHF algorithms except PPO and DPO, such as IPO, KTO, etc. It would be great to see experiments on these algorithms.
2. It would be better to perform experiments on various types of dataset to further verify the results of this paper, as different types of dataset might have distinct properties. For example, it would be interesting to see the results on reasoning dataset.

questions:
Nothing necessary stands out.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730084680126

thread_type:
official_review

<thread_object>
creator: reviewer_7cqn

<content>
summary:
This paper explores the possibility of asynchronous RLHF training for LLMs. Specifically, the generation and training steps, which constitute the majority of RLHF's processing time, are conducted asynchronously on separate GPUs. This approach naturally introduces staleness or so-called "off-policyness," meaning the training samples are generated by earlier versions of the LLM rather than the most recent one.

The authors find that off-policyness can significantly impact the performance of PPO across a wide range of model sizes. In contrast, DPO demonstrates greater resilience to off-policyness. Accordingly, the authors have built a system that employs online asynchronous DPO for RLHF, which can achieve 20% faster training than the synchronous approach in wall-clock time while delivering comparable or superior performance.

Currently, I don't find any critical flaws in this paper and would recommend for acceptance.

strengths:
* This paper presents a novel perspective on off-policyness in RLHF. I especially appreciate the experiments comparing various RLHF loss functions and scaling behaviors related to off-policyness in Section 3. These experiments not only enrich this paper but also validate concurrent works on applying DPO to off-policy data and advocating for on-policy PPO training.

* The paper is well-organized and easy to follow.

weaknesses:
While I don't have significant concerns that would prevent a positive recommendation, it would be helpful if the authors could clarify the following points to avoid potential misunderstandings.

## Motivation

The motivation presented in the first paragraph of Section 3 suggests that training libraries are slow for generation, and using separate libraries for training and generation could lead to GPU underutilization. This appears to be more of an implementation issue than a scientific one, which may not fully justify the exploration of off-policy RLHF. If RLHF systems can indeed achieve sufficient speed, investigating off-policy algorithms may not be necessary. For similar reasons, I have concerns regarding the experiments in Section 3.5.

## Experiments

The title "large-scale asynchronous RLHF" seems overstated. LLaMA-8b is not generally considered large-scale. A more accurate title could be "Asynchronous RLHF with LLaMA-8B."

Figure 9 should also include a comparison with online PPO and in more benchmark tasks. It's good to use TL;DR for analysis in Sec.3, but more diverse evaluation is required for the final experiment. The authors can consider benchmarks like SafeRLHF or MATH.

questions:
+ How can we determine when the system is generation-bound or training-bound? Since generation length may vary during training, the system could be generation-bound at the start but training-bound later. The optimizations in Section 4 require selecting the paradigm in advance, which may not be feasible.

+ It was previously observed that off-policyness negatively affects learning performance in the OpenAI Five paper [1]. A dedicated section explaining this relationship could add clarity.

[1] Dota 2 with Large Scale Deep Reinforcement Learning, Figure 5, https://arxiv.org/pdf/1912.06680
</content>
</thread_object>

</thread>
<thread>
creation date:
1730547184728

thread_type:
official_review

<thread_object>
creator: reviewer_6pcs

<content>
summary:
This paper discusses the possibility to achieve faster and more efficient RLHF training of LLMs by exploiting online but off-policy data. The authors propose to optimize compute efficiency by separating the generation and training tasks in RLHF on different sets of GPUs and executing them asynchronously. However, this requires training on off-policy samples. This paper further explores properties of asynchronous off-policy RLHF on TLDR Summarization benchmark. The results show that the online DPO method is the most robust to off-policy data, and robustness increases with the scale of the policy model. Moreover, some optimization methods targeting generation-bound and training-bound scenarios are also proposed to further increase the efficiency of asynchronous RLHF. Finally, the authors train a LLaMA 3.1 8B model on an instruction following task with online DPO, showing that the asynchronous method is 40% faster than its synchronous counterpart.

strengths:
1. The paper is well-organized and the ideas could be easily understood.
2. The asynchronous training in RLHF of LLMs is an relatively unexplored topic compared to synchronous methods. This paper has made a first attempt to research in this area.
3. The experiments on the robustness of different RLHF algorithms to off-policy data is abundant. It could be inspiring to researches in related topics.

weaknesses:
1. For asynchronous RLHF, it is unclear to me where the computation efficiency increase comes from. Compared to synchronous RLHF, does asynchronous RLHF reduce communication overheads or GPU idle time? Or does it decrease the overall computation required to achieve the same goal? It would be much more easier for readers to understand the performance gain if the authors could provide a breakdown analysis in the experiments. 
2. From the first paragraph in Section 3, it seems that a major obstacle that prevents the synchronous RLHF to be efficient is that there are no efficient implementation of generation in existing training libraries. From my perspective, this is only an issue that could be solved by engineering efforts. If this issue is solved, what is the advantage of asynchronous RLHF? Alternatively, could you explain why it is infeasible to make generation efficient in the training libraries? 
3. The authors have also mentioned in the footnote that there are a more advanced approach to interleave training and generation [1]. However, besides the paged memory in vLLM (which will reserve a large portion of GPU memory for KV cache), there are also many optimizations that could be applied to the generation (such as CUDAGraph, speculative decoding, continuous batching, etc.) in RLHF without harming the performance of training tasks, making it much more faster than the huggingface transformers implementation. If these optimization methods are applied, what is the advantage of asynchronous RLHF compared to the approach in [1]?
4. In section 5, the experiment result only shows the asynchronous version of online DPO compared to the synchronized version of online DPO. However, it is still unclear how the asynchronized online DPO performs compared to other synchronized algorithms with the same amount of computational resources. Could you provide cases in which the asynchronized DPO outperforms the synchronized version of another RLHF algorithm (e.g. PPO)? This would make the advantage of asynchronous RLHF much more convincing.
5. The benchmark used in this paper is TLDR Summarization, which is relatively simple compared to problems such as coding and math.  As an extension to problem 4, will the asynchronized online DPO perform better than the synchronized PPO in these harder problems? 

[1] Mei, Z., Fu, W., Li, K., Wang, G., Zhang, H., & Wu, Y. (2024). ReaLHF: Optimized RLHF Training for Large Language Models through Parameter Reallocation. https://arxiv.org/abs/2406.14088

questions:
Please respond to the questions in the weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730626521649

thread_type:
official_review

<thread_object>
creator: reviewer_dsgt

<content>
summary:
The paper propose a asynchronous approach for accelerating the training of RLHF when multiple GPUs are available. Specifically, they propose to do generation and training simultaneously by letting policy learn on off-policy data. Experiments on TLDR dataset shows that the algorithm can acclerate the training by about 40%. Additionally, the paper provides some empirical results on the effect of how the off-policy data may affect the policy learning.

strengths:
* The asynchronous implementation enables model to do generation and training simultaneously and reduces the time for waiting reward signal. Experiments in finetuning Llama 3.1 shows a 40% acceleration in terms of reward convergence.

weaknesses:
* The study on the effect of using off-policy data for updating is not very convincing. Based on the protocol in the paper (line 200-207), the policy seems to be trained on the data generated by the *same* outdated policy, which may be the key that results in performance drop. Specifically, the data generated by the same policy may be homogeneous (think about an extreme case where the policy generates $N$ same data pairs). In this case, the model only needs to learn a small update to adapt on these data, resulting a small KL divergence. A better experimental design is to train model using data generated by *different model* (that is $N$ steps behind) for each update.

* Given different computational facility and tasks, it is not trivial to design an optimal resource allocation scheme for generation and training, which may require heavy manual tuning.

questions:
Please refer to my weakness part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733205469777

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Discussion Period Summary

comment:
After useful and fruitful discussions with all reviewers, we seem to have assuaged all concerns 
- reviewer 7CQN has no remaining concerns and recommends acceptance noting it is an "excellent paper" and they "will advocate for its acceptance"
- reviewer rLXT has no remaining concerns and recommends acceptance noting it is an "important problem" and "this paper gives a thorough study of [asynchronous RLHF] with meaning[ful] conclusions."
- reviewer 6pCS had a few remaining points that we clarified and recommends acceptance noting it "is a very promising research topic and this paper is beneficial to the community"
- reviewer dsGt had one last minor point to clarify and we hope we have adequately addressed it

We thank all reviewers for their time and effort and note the biggest changes from the discussion period.

**improved motivation**

Thanks to deep discussions with 6pCS and 7CQN, we have greatly strengthened our motivation and added a substantial appendix C describing the engineering details and a case study for the challenges of synchronous RLHF. We go into detail why asynchronous RLHF is a simpler, effective approach. We thank the reviewers for helping clarify our work.

**strong new results in reasoning / math**

Thanks to a suggestion by 6pCS, 7CQN, and rLXT, we have also gone beyond the scope of our initial paper and achieved asynchronous RL with LLMs for math/reasoning. In a new appendix D.2, we compare to a state-of-the-art PPO baseline using the standard GSM8k dataset. Our synchronous Online DPO setup runs nearly 4x faster and achieves better final results than PPO. Our asynchronous Online DPO is comparable and even slightly better than synchronous, while running 68% faster. This speedup is even larger than our previous results and demonstrates the promise of asynchronous RL for LLMs. We thank the reviewers for the suggestion and we plan to incorporate this very strong result into the main paper.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper discusses the possibility to achieve faster and more efficient RLHF training of LLMs by exploiting online but off-policy data. The authors propose to optimize compute efficiency by separating the generation and training tasks in RLHF on different sets of GPUs and executing them asynchronously. However, this requires training on off-policy samples. This paper further explores properties of asynchronous off-policy RLHF on TLDR Summarization benchmark. The results show that the online DPO method is the most robust to off-policy data, and robustness increases with the scale of the policy model. Moreover, some optimization methods targeting generation-bound and training-bound scenarios are also proposed to further increase the efficiency of asynchronous RLHF. Finally, the authors train a LLaMA 3.1 8B model on an instruction following task with online DPO, showing that the asynchronous method is 40% faster than its synchronous counterpart.

All reviewers and the AC agree that this paper studies an emerging and important problem and this paper makes non-trivial technical contributions. The AC thus recommends acceptance.