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

We thank the reviewer for their comments and are glad our problem is important, our results are promising, and our conclusions are meaningful.

Other *PO Methods

• PPO and Online DPO are current SOTA methods see 1 and 2
• IPO has shown very similar results to DPO in both the original paper, and follow-up work, we expect similar results here
• KTO is a method for un-paired data. Our problems (summarization, general-purpose chatbots) generally have paired comparison data.

We believe our results demonstrate that asynchronous versions of other *PO methods would also be faster and more efficient

Another Benchmark

On top of TLDR summarization and No Robots general-purpose chatbot, we are currently running experiments with GSM 8k as a reasoning / math benchmark. Thank you for the suggestion.

We thank the author for their review and appreciate that they recommend our work be accepted and find it well-organized, novel, and even enriching other concurrent work.

Why is efficient generation infeasible for synchronous RLHF

We thank the reviewer for the question. We answer below and have added a new Appendix C with a in-depth explanation.

We use SOTA synchronous RLHF library, NVIDIA’s NeMo-Aligner, as an example. Training and Generation must use fundamentally different backends for efficiency (e.g. they use Megatron-LM for training, TensorRT-LLM for generation). NeMo-Aligner is an impressive engineering effort that integrates them both in synchronous RLHF but it adds overhead and still isn’t feasible to maintain because generation and training libraries are advancing too quickly in opposite directions.

To make it work, NeMo-Aligner’s team wrote complex code to convert between training (Megatron-LM) and generation (TensorRT-LLM) backends. But maintaining this is difficult! Once the generation backend (TensorRT-LLM) was updated to v12 and then v13, it took the NeMO-Aligner team 1.5 months to integrate it into their library. By that time, TensorRT-LLM was updated again to a new version v14 and NeMO-Aligner is again out-of-date. And this is a best-case scenario since all these libraries are from NVIDIA.

Synchronous training libraries must consistently make big changes to integrate the latest generation code. Asynchronous training treats generation like an API and integrates new code easily.

Section 5 Title

We have changed the title of Section 5, thank you

PPO for 8B Chatbot

We are currently running synchronous and asynchronous PPO for our 8B general-purpose chatbot. Thank you for the suggestion.

Another Benchmark

On top of TLDR summarization and No Robots general-purpose chatbot, we are currently running experiments with GSM 8k as a reasoning / math benchmark. Thank you for the suggestion.

Training vs Generation-Bound

In order to account for compute changes during training, it is possible to make best-of-$K$ samping and num ppo epochs $N$ into hyperparameters that anneal / change with time or compute considerations. In practice, we didn't find significant changes during our training.

Off-Policy Degradation in OpenAI Five

This is a great pointer, thank you. We’ve added this citation and note to the end of Section 3.2

Thank you for your quick response and deep dive into our problem

1. Determining training/generation-bound

We keep things simple and launch a short synchronous test experiment to get an idea of training/generation speeds per iteration. We then determine how many GPUs to use and divide our GPUs using the ratio of training/generation speeds.

2a. Our research contribution vs engineering work (NeMo-Aligner)

We’d like to clarify a major point: we are not proposing an engineering solution but a research paradigm (asynchronous + off-policy learning). Our contribution is not directly comparable to engineering works such as NeMo-Aligner or HybridFlow. These libraries show how to optimize allocations and pipelines in the existing paradigm, but they don’t suggest new paradigms, they don’t show learning curves, or convergence. They focus purely on throughput of tokens. E.g. NeMo-Tensor’s LLaMA-3 model trained with PPO is worse than the LLaMA-3 instruct model and HybridFlow doesn’t show a single learning curve or demonstrate any results that its pipeline can effectively train a model.

In contrast, we demonstrate that asynchronous learning is a new paradigm that can achieve equal performance more efficiently, an alternative to the synchronous status quo. We focus on the effectiveness of the trained model. Since this paradigm is novel, there does not exist a fully optimized library for asynchronous RLHF for comparison, and creating it would be a big engineering effort.

2b. Our Motivation is General: HybridFlow is already out of date

Integrating updates is not just an issue for NeMo-Tensor, but all synchronous RLHF libraries that aim to integrate training and generation libraries into one. The reviewer’s suggestion, HybridFlow, has it much worse than NeMo-Tensor. Despite being released 3 weeks ago, it uses Megatron 0.4.0 and vllm 0.5.4. These libraries are, respectively, 11 months and 3.5 months out of date. HybridFlow, despite being just released, does not support MoE, LLaMA 3.2, or multistep scheduling which 2x improves generation throughput.

We believe our motivation is a fundamental limitation: not fundamental scientific possibility but fundamental feasibility. Whereas synchronous libraries have complex combinations of tricks to co-locate training and generation, our proposed asynchronous paradigm is effective while being drastically simpler.

We thank the reviewer for their in-depth assessment and critiques, we appreciate that they find our paper well-organized, novel, and inspiring.

Here, we try to clarify the why of asynchronous learning and we’ve added a new Appendix C with a more in-depth explanation.

1. Why is async more efficient

Our speed improvements come from SOTA generation libraries, which have separate optimizations from training. In Appendix C.1 we show how much faster vllm is than huggingface transformers.

Synchronous RLHF is slower since either (1) it uses training libraries for generation, which makes generation really slow or (2) it uses SOTA generation on separate GPUs, but incurs GPU idle time: training GPUs are idle while waiting for generation GPUs to finish and vice versa. Asynchronous RLHF uses SOTA generation on separate GPUs and reduces GPU idle time by parallelizing generation and training.

2. Why is efficient generation infeasible for synchronous RLHF

Syncronous RLHF can also (3) interleave SOTA generation and training on the same GPUs. This is actually available in NVIDIA’s SOTA RLHF library, NeMo-Aligner, an impressive engineering effort. But it adds communication overhead and we argue it isn’t feasible to maintain as generation and training libraries are advancing too quickly in opposite directions.

To make it work, NeMo-Aligner’s team wrote complex code to convert between the separate training (Megatron-LM) and generation (TensorRT-LLM) libraries. But maintaining this is difficult! Once the generation library (TensorRT-LLM) was updated to v12 and then v13, it took the NeMO-Aligner team 1.5 months to integrate it into their library. By that time, TensorRT-LLM was updated again to a new version v14 and NeMO-Aligner is again out-of-date. And this is a best-case scenario since all these libraries are from NVIDIA.

Synchronous training libraries must consistently make big changes to integrate the latest generation code. Asynchronous training treats generation like an API and integrates new code easily. Thank you for this question and please see our new Appendix C.2 for a more detailed explanation.

3. Why can’t training use generation optimizations

Many newer optimizations are incompatible with training (e.g. paged-attention). Other generation tricks can be integrated in synchronous training but you usually pay a cost. E.g. NeMo-Aligner compiles the model into a CUDA graph but this takes up memory on the GPU. This means less available memory during training so NeMo-Aligner has to compensate with even more tricks to reduce optimizer memory (gradient accumulation and gradient checkpointing) but these make training slower.

4. Sync vs Async PPO

We do not necessarily recommend async Online DPO, but rather async vs sync. We are currently running sync and async PPO for our general-purpose 8B chatbot to show that async PPO can match sync PPO while being faster. Thank you for the suggestion.

5. Another benchmark

On top of TLDR summarization and No Robots general-purpose chatbot, we are currently running experiments with GSM 8k as a reasoning / math benchmark. Thank you for the suggestion.

Thank you for the detailed responses and the addtional appendix. However, a few concerns still remain:

1. In Appendix C.2., you have mentioned the overheads in dynamic model resharding. It would be more convincing if you could give a small example on how much the overhead is in a typical training experiment using NeMo-Aligner. Specifically, how much percentage of time does it take to reshard parameters in one training iteration?

2. You have mentioned that NeMo-Aligner's team wrote complex code to convert training and generation libraries, and it is difficult to maintain the code when the libraries update. This seems a specific engineering flaw in the implementation of NVIDIA's training and generation library. I understand that the authors may not know about the underlying details of the development of NeMo-Aligner and related NVIDIA's libraries, but it seems that an important argument in the paper relies on this specific engineering flaw and it is not well-discussed.

3. About CUDAGraph, it does take extra memory on the GPUs and could not be freed during training in NeMo-Aligner. However, it is not necessary since CUDAGraph could be destroyed before training and recaptured before the next generation. In this way, the memory space for CUDAGraph will be freed when training. I believe that this is another engineering flaw in NeMo-Aligner.

To show that your argument is not solely established on the engineering flaws of NeMo-Aligner, I recommend that you could provide an end-to-end comparison between your asynchronous implementation and synchronous open-source libraries including NeMo-Aligner and others (e.g. DSChat).

Thank you for the fast response and reading our appendix, we appreciate this discussion.

1. Resharding Overhead

We are not experts in NeMo-Aligner, so we can’t give exact estimates and overhead as a percentage depends on the number of tokens generated for responses, number of PPO epochs etc.. Since this resharding is a dynamic design decision mentioned by the authors, we believe it is time-consuming enough to matter.

2. Training and Generation must use different backends

Every state-of-the-art RLHF library uses different backends for training and generation to better optimize each. Training is a dynamic process which focuses on faster backprop and generation is a fixed process that focuses on faster throughput.

• RealHF (Megatron-LM/Deepspeed and vLLM)
• NeMo-Tensor (Megatron-LM, TensorRT-LLM)
• OpenRLHF (Deepspeed, vLLM)

It is not a flaw, but an achievement that the backends are different: it shows how much generation has been optimized in ways that don’t apply to training. RLHF libraries that try to use one backend are too slow e.g. OpenRLHF notes “the generation acceleration of vLLM is substantially better than that of [Deepspeed-Chat] Hybrid Engine”.

We hope this is convincing that we are motivated by a fundamental limitation of all synchronous RLHF libraries

3. CUDA graph / engine cannot be freed and recreated

Building the CUDA graph is very computationally expensive and prohibitive to do each iteration. From NeMo-Aligner, “due to the cost of serializing the engine, we keep the engine in memory during training” and “the first step incurs substantial time for TensorRT-LLM Engine Building”

Our research contribution vs engineering work

We’d like to clarify a major point: we are not proposing an engineering solution but a research paradigm (asynchronous + off-policy learning). Our contribution is not directly comparable to NeMo-Aligner, an engineering work. Engineering works optimize allocations and pipelines in an existing paradigm and focus purely on throughput of tokens. They don’t suggest new paradigms, show learning curves, or look at convergence.. E.g. NeMo-Tensor’s LLaMA-3 model trained with PPO is worse than the LLaMA-3 instruct model. Comparison to NeMo-Aligner was also not possible at submission time since, as noted in Appendix C.3, it did not support LLaMA 3.1

In contrast, we demonstrate that asynchronous learning is a new paradigm that can train models with equal performance but more efficiently, an alternative to the synchronous status quo. We focus on the effectiveness of the trained model. Since this paradigm is novel, there does not exist a fully optimized library for asynchronous RLHF, which would be a big engineering effort.

Thank you for your notes and continued discussion, it is helping us improve our work even further. We also have novel results to share.

1. This work is about the viability and promise of Asynchronous RLHF

We agree that our asynchronous RLHF is slower than current state-of-the-art synchronous RLHF libraries: changing a paradigm and outperforming previous SOTA is a large task, arguably a whole research agenda. Lots of effort has gone into optimizing synchronous RLHF which makes it hard to immediately compete with. But to avoid getting stuck in a hardware lottery, we have demonstrated the viability and advantages of asynchronous RLHF. Our work is the beginning of the research agenda, laying the groundwork so the community can investigate and begin to optimize this different paradigm.

We have softened the conclusion of our paper in line with this idea, thank you.

2. Why Asynchronous is more efficient

In Appendix C.1 we demonstrate the need for generation libraries. In section 5.2 and at the end of the new Appendix D (on the final page), we examine the exact runtimes and how asynchronous learning speeds up training over a simple, but competitive, synchronous RLHF by reducing idling time.

3. Difficult maintenance is a fact of all synchronous, integrated RLHF libraries

Synchronous RLHF libraries manually integrate training and generation libraries to work around synchronous limitations (constant switching between train and generate). Updates generally break those manual integrations.

NeMo-Aligner is not just an example but a best case scenario. It is maintained by a team of strong engineers and the underlying training and generation frameworks are from the same company. For an average case, HybridFlow was released 3 weeks ago, it uses Megatron 0.4.0 and vllm 0.5.4. These libraries are, respectively, 11 months and 3.5 months out of date. HybridFlow, despite being just released, does not support MoE, LLaMA 3.2, or multistep scheduling which 2x improves generation throughput.

Novel Result: Sync vs Async PPO experiments

We have run Sync vs Async PPO for our 8B chatbot, please see the new Appendix D.1. We find that, once again, async training matches sync while being faster. Online DPO is still more performant than PPO for this task. Thank you for this suggestion, we believe it strengthens our contribution.

Novel Result: Asynchronous RL for Math/Reasoning on GSM8k

We have gone beyond the scope of original work and achieved strong results with asynchronous training on the math and reasoning benchmark GSM8k. In Appendix D.2 we show how synchronous Online DPO improves on performance compared to an existing, highly-tuned PPO baseline. We then show how asynchronous Online DPO achieves the same final performance as synchronous with an even bigger speedup than any previous result, 68% faster. Overall, we outperform a state-of-the-art synchronous PPO baseline while running 6x faster on comparable hardware.

This is a big result and we thank the reviewer for suggesting it. We believe we should integrate it into the main results of the paper, replacing section 5.2 and extend our conclusion to generally RL for language models. We are open to feedback and suggestions.

We thank the reviewer for their comments and we’re glad they appreciate our improvements in speed and efficiency. We agree this is a key part of our paradigm.

Off-Policy Setup: Same vs Different Policy?

Our setup is made to mimic how asynchronous learning requires off-policy learning. For the off-policy setup, we already sample from our own policy. And as we update, our final sample is from our own policy, now $N$ steps behind. Empirically, we don’t see any issue of homogenous data since we sample with temperature 1 and there is a consistent difference in reward between our two responses per prompt.

Other work has already noted that off-policyness should reduce performance. Can the reviewer explain which exact model is being suggested for sampling?

Resource Allocation

We agree that designing an optimal resource allocation strategy can be difficult. However, our tuning is straightforward: if generation is slower then use more GPUs for generation, and if training is slower we can run more GPUs for training. And synchronous RLHF libraries like NeMO-Aligner already deal with this because they use asynchronous reward models and have existing ways for reward and policy compute allocation.

We thank the reviewer for the comment and have added a new section in Appendix C.3 to discuss the issue of resource allocation. Does this alleviate the concern?

Thanks for the response, but my concern is not well addressed.

Specifically, consider the scenario where generation is slower than training. After training model for the first $N$ steps ($N > 1$), there is less than $N$ samples generated. Then, where does the next data batch come from? Could authors explain the detailed protocol? These information are not included in Algorithm.

Our Cleanba (Huang et al, 2023) style asynchronous learning is not completely asynchronous. Training and generation sync every $N$ steps (e.g. $N=1$ in our general-purpose chatbot experiments in Section 5). So if generation is slower, training will idle for generation to finish (and vice-versa). We will make this detail clearer in the Algorithm in the Appendix.

As shown in Figure 6, training will wait for generation to finish in order to get the next batch of data. Asynchronous RLHF does not eliminate idling, but can reduce it greatly.