trl X : A Framework for Large Scale Reinforcement Learning from Human Feedback (2023)

Abstract: Reinforcement learning from human feedback ( RLHF ) utilizes human feedback to better align large language models with human preferences via online optimization against a learned reward model. Current RLHF paradigms rely on Proximal Policy Optimization ( PPO ), which quickly becomes a challenge to implement and scale up to large architectures. To address this difficulty we present the AutoRLHF library as a feature complete open-source framework for RLHF fine-tuning of models up to and exceeding 70 billion parameters. To do so we implement support for multiple types of distributed training including distributed data parallel, model sharded, as well as tensor, sequential, and pipeline parallelism. Additionally, we implement compute and memory saving features, giving AutoRLHF the flexibility to support users with a wide range of compute resources. This includes offline RL methods like Implicit Language Q Learning ( ILQL ) as a compute efficient alternative to PPO. We find offline fine-tuning offers competitive performance relative to online algorithms while being easier to implement, train, and scale. To evaluate our framework we train RLHF models on two separate well-known tasks using publicly available human preference data. Models trained with AutoRLHF achieve preference win-rates over baselines at rates comparable to the original works.

## introduction since 2019, the prevailing training paradigm for large-language models @xcite @xcite has comprised two parts: first a pre-training phase during which models are trained on a large corpus of text. pretraining produces a general-purpose model which has learned the syntax and some semantics of natural text but is not easily controllable for any particular task. fine-tuning is then used to adapt this generalpurpose pre-trained model to a smaller, specialized corpus. however, this often comes with a decrease in model performance on other tasks. further, the resulting model can be difficult to work with, requiring extensive prompt engineering. reinforcement learning from human feedback (rlhf), has presented itself as an additional third stage of the language model training pipeline. in this stage, gathered human preference data is used to supervise fine-tune the pre-trained model and then train a reward model. the reward model assigns scalar values to (prompt, response) pairs that correspond to human preference. the supervised fine-tuned model can then be trained using proximal policy optimization (ppo) @xcite , an online reinforcement learning algorithm, to optimize against the learned reward model. the resulting models are better aligned with human preferences, leading to win-rates over the base model by up to 80% @xcite @xcite , and producing more desirable and less harmful text with less prompting. however, the study of rlhf models in the academic community has been limited thus far by a need for open-source training frameworks and datasets. online reinforcement learning via ppo is compute expensive and challenging to scale as we must store three copies of the model in memory during training. to address these issues we present trlx: a library supporting online and offline rlhf fine-tuning for language models of up to and exceeding 70 billion parameters. at smaller scales, we emphasize low-resource accessibility by incorporating features such as hydra model architectures @xcite , lora adapters @xcite , and deepspeed @xcite , which, when combined, can reduce memory overhead on gpt-j @xcite by up to 75% with minimal impact on the achieved reward. trlx training at this scale is compatible with most encoder-decoder and decoder-only architecture supported on the popular hugging face hub @xcite . for large, 20 billion parameter plus, training jobs we implement support for tensor, sequence, and pipeline parallelism via both the gpt-neox library @xcite and nemo-megatron @xcite . trlx also supports implicit language q learning, ilql, as an offline alternative to online rl methods. we find fine-tuning via ilql achieves preference win-rates with baselines close to ppo but at a fraction of the compute cost. further ilql is more robust to reward model overfitting, which online algorithms can suffer from. to evaluate our framework we provide open-source replications of well-known papers from the rlhf literature including learning summaries from human feedback @xcite and helpful/harmless preference learning for a general purpose language assistant @xcite . we find models trained with trlx achieve preference win-rates, as judged by human annotators, over baselines at rates comparable to the original works, validating our implementation. we open-source all supervised fine-tuned models, reward models, and rlhf models for further research, as well as the training framework, trlx. in summary, we make the following contributions: reinforcement learning from human feedback reinforcement learning from human feedback attempts to improve agent performance, either in sample efficiency or performance on downstream tasks, by incorporating some form of human feedback on agent behavior @xcite @xcite . we focus on the 3-stage fine-tuning pipeline outlined by @xcite in which researchers first assemble a dataset of human-annotated preferences, then train a reward model to predict these preferences, and lastly train a policy to maximize the score of the resulting reward model. similarly structured pipelines have been adapted to train many of the most recent interactive natural language assistance tools @xcite @xcite . nevertheless, collecting human preferences at scale can be cost-inefficient. recent work by @xcite proposes tractable oversight by using synthetic ai preferences instead of direct human labels. similar works by @xcite ; @xcite generate instructionfollowing datasets by querying already aligned models like text-davinci-003 to generate both task and instruction-following responses. scalable training frameworks there are many notable choices of distributed training frameworks for large-scale language model pre-training and fine-tuning, each implementing various parallelism schemes. these include deepspeed @xcite , megatron-lm @xcite , the gpt-neox library @xcite which combines deepspeed and megatron-lm, fairseq @xcite , and t5x @xcite for tpu-based training. however, none of these frameworks are designed explicitly to support fine-tuning via rl and would, therefore, require significant work to retrofit. more recently, rl-specific fine-tuning libraries have become available. rl4lm's @xcite implements online algorithms for fine-tuning mid-sized language models with reinforcement learning from human feedback and supports an impressive range of tasks and metrics. trl @xcite , initially a smaller-scale library for transformer reinforcement learning, is a re-implementation of @xcite in py-torch for doing sentiment-based fine-tuning. it has since been expanded in a manner similar to trlx to include training with deepspeed via hugging face accelerate. more recently, deepspeed-chat @xcite was released allowing for larger model training with better throughput. however, they do not allow for more advanced forms of parallelism supported in trlx. the framework is built around a base trainer from which integration specific trainers can inherit. independently, online and offline algorithms are implemented allowing for reuse in separate integrations. in particular trlx supports ppo and a2c for online rl and ilql for offline rl. the most expensive part of online ppo training is the model rollout, which can take up to 10x as long as a combined forward and backward pass. to efficiently maximize batch size for both rollouts and optimization steps, we decouple the rollout inference batch size from the ppo batch size via an orchestrator class. this allows online models to perform batched rollouts to reduce the amount of bottleneck time spent infering each model. we integrate closely with the hugging face ecosystem, allowing for the training of most encoder-decoder and decoder-only models on the hugging face hub, including widely used models such as t5 @xcite and flan-t5 @xcite , gpt-j @xcite , pythia @xcite , opt @xcite , and llama @xcite . fine-tuning large language models via human feedback with ppo is prohibitively expensive in terms of memory and flops, requiring the user to store a student model, reference model, and similarly sized reward model in memory at all times. additionally, reinforcement learning is notoriously brittle to hyperparameter choice, often requiring extensive sweeping to find optimal settings. to mitigate these costs we support parameter saving techniques like lora @xcite and a hydra model architecture design @xcite which allows for shared frozen layers between policy, value, and reference networks. similarly, ilql models require non-standard q-value heads and generation capabilities which are implemented separately for both integrations. ## benchmarks and results we benchmark trlx on two popular rlhf tasks: openai's learning to summarize tl;dr dataset @xcite ) and anthropic's helpful qa dataset @xcite . we release all associated code and models as open-source for further study. training setup and hyperparameters unless otherwise noted, we use the same fixed set of hyperparameters listed in the appendix across all training runs. we find good performance is particularly sensitive to some parameters: ## conclusion ethics we present trlx as an open-source framework for training large language models at scale using reinforcement learning from human feedback. even with fine-tuning via rlhf, llms are still prone to hallucination and bias at inference time, necessitating further research on mitigations. we hope researcher access to trlx as an open-source implementation of the rlhf pipeline will help facilitate this research. limitations while it is highly performant, ppo suffers from many limitations both in difficulty during implementation and in hyperparameter sensitivity at train time. offline methods such as ilql are both easier to implement and more compute efficient, but still do not achieve the same performance as ppo.