Improving your LLMs with RLHF on Amazon SageMaker

Reinforcement Learning from Human Feedback (RLHF) is recognized as the industry standard technique for ensuring large language models (LLMs) produce content that is truthful, harmless, and helpful. The technique operates by training a “reward model” based on human feedback and uses this model as a reward function to optimize an agent’s policy through reinforcement learning (RL). RLHF has proven to be essential to produce LLMs such as OpenAI’s ChatGPT and Anthropic’s Claude that are aligned with human objectives. Gone are the days when you need unnatural prompt engineering to get base models, such as GPT-3, to solve your tasks.

An important caveat of RLHF is that it is a complex and often unstable procedure. As a method, RLHF requires that you must first train a reward model that reflects human preferences. Then, the LLM must be fine-tuned to maximize the reward model’s estimated reward without drifting too far from the original model. In this post, we will demonstrate how to fine-tune a base model with RLHF on Amazon SageMaker. We also show you how to perform human evaluation to quantify the improvements of the resulting model.

Prerequisites

Before you get started, make sure you understand how to use the following resources:

• Amazon SageMaker notebook instances
• Use Amazon SageMaker Ground Truth to Label Data

Solution overview

Many Generative AI applications are initiated with base LLMs, such as GPT-3, that were trained on massive amounts of text data and are generally available to the public. Base LLMs are, by default, prone to generating text in a fashion that is unpredictable and sometimes harmful as a result of not knowing how to follow instructions. For example, given the prompt, “write an email to my parents that wishes them a happy anniversary”, a base model might generate a response that resembles the autocompletion of the prompt (e.g. “and many more years of love together”) rather than following the prompt as an explicit instruction (e.g. a written email). This occurs because the model is trained to predict the next token. To improve the base model’s instruction-following ability, human data annotators are tasked with authoring responses to various prompts. The collected responses (often referred to as demonstration data) are used in a process called supervised fine-tuning (SFT). RLHF further refines and aligns the model’s behavior with human preferences. In this blog post, we ask annotators to rank model outputs based on specific parameters, such as helpfulness, truthfulness, and harmlessness. The resulting preference data is used to train a reward model which in turn is used by a reinforcement learning algorithm called Proximal Policy Optimization (PPO) to train the supervised fine-tuned model. Reward models and reinforcement learning are applied iteratively with human-in-the-loop feedback.

The following diagram illustrates this architecture.

In this blog post, we illustrate how RLHF can be performed on Amazon SageMaker by conducting an experiment with the popular, open-sourced RLHF repo Trlx. Through our experiment, we demonstrate how RLHF can be used to increase the helpfulness or harmlessness of a large language model using the publicly available Helpfulness and Harmlessness (HH) dataset provided by Anthropic. Using this dataset, we conduct our experiment with Amazon SageMaker Studio notebook that is running on an ml.p4d.24xlarge instance. Finally, we provide a Jupyter notebook to replicate our experiments.

Complete the following steps in the notebook to download and install the prerequisites:

git clone https://github.com/CarperAI/trlx.git
cd trlx
pip install torch==2.0.0 --extra-index-url https://download.pytorch.org/whl/cu116 # for cuda
pip install -e .

Import demonstration data

The first step in RLHF involves collecting demonstration data to fine-tune a base LLM. For the purpose of this blog post, we’re using demonstration data in the HH dataset as reported above. We can load the demonstration data directly from the Hugging Face datasets package:

from datasets import load_dataset
dataset = load_dataset("Dahoas/rm-static")

Supervised fine-tuning a base LLM

The next step is to perform supervised fine-tuning of a base LLM. In this blog post, we refer to the base model that has undergone supervised fine-tuning simply as the “SFT model”. Supervised fine-tuning is required to learn from demonstration data so that an LLM performs well on our conversational task and learns to be helpful and harmless. In this post, we use the publicly available EleutherAI/gpt-j-6b model hosted on Hugging Face. We also use the Trlx framework that provides code for supervised fine-tuning this model.

Run the following commands to begin training:

cd examples/hh
accelerate launch --num_processes 7 --config_file ../../configs/accelerate/zero2-bf16.yaml sft_hh.py

Import preference data

As shown in the previous diagram, a critical step in RLHF involves acquiring preference data. Preference data is a collection of examples that demonstrate how a human prefers one machine output over another based on helpfulness and harmlessness criteria.

The following table shows the concept of preference:

Train your reward model

Our reward model is based on GPT-J-6B and is fine-tuned on the previously mentioned HH dataset. Since training the reward model is not the focus of this post, we will use a pre-trained reward model specified in the Trlx repo, the Dahoas/gptj-rm-static. If you want to train your own reward model, please refer to the autocrit library on GitHub.

RLHF Training

Now that we have acquired all the required components for RLHF training (i.e., an SFT model and a reward model), we can now begin optimizing the policy using RLHF.

To do this, we modify the path to the SFT model in examples/hh/ppo_hh.py:

elif config_name == "6B":
    ...
    default_config.model.model_path = PATH_TO_THE_SFT_MODEL_IN_THE_PREVIOUS_STEP
    ...

We then run the training commands:

cd examples/hh 
CONFIG_NAME=6B accelerate launch --num_processes 7 --config_file ../../configs/accelerate/zero2-bf16.yaml ppo_hh.py

The script initiates the SFT model using its current weights and then optimizes them under the guidance of a reward model, so that the resulting RLHF trained model aligns with human preference. The following diagram shows the reward scores of model outputs as the RLHF training progresses. Reinforcement training is highly volatile, so the curve fluctuates, but the overall trend of the reward is upward, meaning that the model output is getting more and more aligned with human preference according to the reward model. Overall, the reward improves from -3.42e-1 at the 0-th iteration to the highest value of -9.869e-3 at the 3000-th iteration.

The following diagram shows an example curve when running RLHF.

Human evaluation

Having fine-tuned our SFT model with RLHF, we now aim to evaluate the impact of the fine-tuning process as it relates to our broader goal of producing responses that are helpful and harmless. In support of this goal, we compare the responses generated by the model fine-tuned with RLHF to responses generated by the SFT model. We experiment with 100 prompts derived from the test set of the HH dataset. We programmatically pass each prompt through both the SFT and the fine-tuned RLHF model to obtain two responses. Finally, we ask human annotators to select the preferred response based on perceived helpfulness and harmlessness.

The Human Evaluation approach is defined, launched, and managed by the Amazon SageMaker Ground Truth Plus labeling service. SageMaker Ground Truth Plus enables customers to prepare high-quality, large-scale training datasets to fine-tune foundation models to perform human-like generative AI tasks. It also allows skilled humans to review model outputs to align them with human preferences. Additionally, it enables application builders to customize models using their industry or company data while preparing training datasets. As shown in a previous blog post (“High-quality human feedback for your generative AI applications from Amazon SageMaker Ground Truth Plus”), SageMaker Ground Truth Plus provides workflows, labeling interfaces, and a skilled workforce to generate high-quality human feedback data to accomplish these tasks in a turnkey, fully managed service on behalf of the customer. We use a task interface similar to the UI shown below to facilitate the human evaluation step of our experiment.

We conclude our model evaluation experiment by determining which of the two models – SFT or RLHF – produced responses that were more frequently preferred. For 84 of the 100 prompts and responses being evaluated, we observe that the responses from the model fine-tuned with RLHF are at least as good as or more preferred than those generated by the model fine-tuned with SFT (RLHF-trained model is more preferred 30 times, tied 54 times). As shown in the table below, a qualitative comparison of model responses provides supporting evidence that illustrates how the RLHF model tends to be more helpful and factually correct than the SFT model. Taken as a collective, our experiment suggests that RLHF does, in fact, improve the model’s ability to produce output that is better aligned with the values of interest.

Toxicity evaluation

To quantify how RLHF reduces toxicity in the model generations, we benchmark on the popular RealToxicityPrompt test set and measure toxicity on a continuous scale from 0 (Not Toxic) to 1 (Toxic). We randomly select 1,000 test cases from the RealToxicityPrompt test set and compare the toxicity of the SFT and RLHF model outputs. Through our evaluation, we find that the RLHF model achieves a lower toxicity (0.129 on average) than SFT model (0.134 on average), which demonstrates the effectiveness of RLHF technique in reducing output harmfulness.

Clean up

Once you’re finished, you should delete the cloud resources that you created to avoid incurring additional fees. If you opted to mirror this experiment in a SageMaker Notebook, you need only halt the notebook instance that you were using. For more information, refer to the AWS Sagemaker Developer Guide’s documentation on “Clean Up”.

Conclusion

In this post, we showed how to train a base model, GPT-J-6B, with RLHF on Amazon SageMaker. We provided code explaining how to fine-tune the base model with supervised training, train the reward model, and RL training with human reference data. We demonstrated that the RLHF trained model is preferred by annotators. Now, you can create powerful models customized for your application.

If you need high-quality training data for your models, such as demonstration data or preference data, Amazon SageMaker can help you by removing the undifferentiated heavy lifting associated with building data labeling applications and managing the labeling workforce. When you have the data, use either the SageMaker Studio Notebook web interface or the notebook provided in the GitHub repository to get your RLHF trained model.

About the Authors

Weifeng Chen is an Applied Scientist in the AWS Human-in-the-loop science team. He develops machine-assisted labeling solutions to help customers obtain drastic speedups in acquiring groundtruth spanning the Computer Vision, Natural Language Processing and Generative AI domain.

Erran Li is the applied science manager at humain-in-the-loop services, AWS AI, Amazon. His research interests are 3D deep learning, and vision and language representation learning. Previously he was a senior scientist at Alexa AI, the head of machine learning at Scale AI and the chief scientist at Pony.ai. Before that, he was with the perception team at Uber ATG and the machine learning platform team at Uber working on machine learning for autonomous driving, machine learning systems and strategic initiatives of AI. He started his career at Bell Labs and was adjunct professor at Columbia University. He co-taught tutorials at ICML’17 and ICCV’19, and co-organized several workshops at NeurIPS, ICML, CVPR, ICCV on machine learning for autonomous driving, 3D vision and robotics, machine learning systems and adversarial machine learning. He has a PhD in computer science at Cornell University. He is an ACM Fellow and IEEE Fellow.

Koushik Kalyanaraman is a Software Development Engineer on the Human-in-the-loop science team at AWS. In his spare time, he plays basketball and spends time with his family.

Xiong Zhou is a Senior Applied Scientist at AWS. He leads the science team for Amazon SageMaker geospatial capabilities. His current area of research includes computer vision and efficient model training. In his spare time, he enjoys running, playing basketball and spending time with his family.

Alex Williams is an applied scientist at AWS AI where he works on problems related to interactive machine intelligence. Before joining Amazon, he was a professor in the Department of Electrical Engineering and Computer Science at the University of Tennessee . He has also held research positions at Microsoft Research, Mozilla Research, and the University of Oxford. He holds a PhD in Computer Science from the University of Waterloo.

Ammar Chinoy is the General Manager/Director for AWS Human-In-The-Loop services. In his spare time, he works on positivereinforcement learning with his three dogs: Waffle, Widget and Walker.