Posts in category: Amazon AWS

This post was co-authored with Daniele Chiappalupi, participant of the AWS student Hackathon team at ETH Zürich.

Everyone can easily get started with machine learning (ML) using Amazon SageMaker JumpStart. In this post, we show you how a university Hackathon team used SageMaker JumpStart to quickly build an application that helps users identify and remove biases.

“Amazon SageMaker was instrumental in our project. It made it easy to deploy and manage a pre-trained instance of Flan, offering us a solid foundation for our application. Its auto scaling feature proved crucial during high-traffic periods, ensuring that our app remained responsive and users received a steady and fast bias analysis. Further, by allowing us to offload the heavy task of querying the Flan model to a managed service, we were able to keep our application lightweight and swift, enhancing user experience across various devices. SageMaker’s features empowered us to maximize our time at the hackathon, allowing us to focus on optimizing our prompts and app rather than managing the model’s performance and infrastructure.”

– Daniele Chiappalupi, participant of the AWS student Hackathon team at ETH Zürich. 

Solution overview

The theme of the Hackathon is to contribute to the UN sustainable goals with AI technology. As shown in the following figure, the application built at the Hackathon contributes to three of the Sustainable Development Goals (quality education, targeting gender-based discrimination, and reduced inequalities) by helping users identify and remove biases from their text in order to promote fair and inclusive language.

As shown in the following screenshot, after you provide the text, the application generates a new version that is free from racial, ethnical, and gender biases. Additionally, it highlights the specific parts of your input text related to each category of bias.

In the architecture shown in the following diagram, users input text in the React-based web app, which triggers Amazon API Gateway, which in turn invokes an AWS Lambda function depending on the bias in the user text. The Lambda function calls the Flan model endpoint in SageMaker JumpStart, which returns the unbiased text result via the same route back to the front-end application.

Application development process

The process of developing this application was iterative and centered on two main areas: user interface and ML model integration.

We chose React for the front-end development due to its flexibility, scalability, and powerful tools for creating interactive user interfaces. Given the nature of our application—processing user input and presenting refined results—React’s component-based architecture proved ideal. With React, we could efficiently build a single-page application that allowed users to submit text and see de-biased results without the need for constant page refreshes.

The text entered by the user needed to be processed by a powerful language model to scrutinize for biases. We chose Flan for its robustness, efficiency, and scalability properties. To utilize Flan, we used SageMaker JumpStart, as shown in the following screenshot. Amazon SageMaker made it easy to deploy and manage a pre-trained instance of Flan, allowing us to focus on optimizing our prompts and queries rather than managing the model’s performance and infrastructure.

Connecting the Flan model to our front-end application required a robust and secure integration, which was achieved using Lambda and API Gateway. With Lambda, we created a serverless function that communicates directly with our SageMaker model. We then used API Gateway to create a secure, scalable, and readily accessible endpoint for our React app to invoke the Lambda function. When a user submitted text, the app triggered a series of API calls to the gateway—first to identify if any bias was present, then, if necessary, additional queries to identify, locate, and neutralize the bias. All these requests were routed through the Lambda function and then to our SageMaker model.

Our final task in the development process was the selection of prompts to query the language model. Here, the CrowS-Pairs dataset played an instrumental role because it provided us with real examples of biased text, which we utilized to fine-tune our requests. We selected the prompts by an iterative process, with the objective of maximizing accuracy in bias detection within this dataset.

Wrapping up the process, we observed a seamless operational flow in the finished application. The process begins with a user submitting text for analysis, which is then sent via a POST request to our secure API Gateway endpoint. This triggers the Lambda function, which communicates with the SageMaker endpoint. Consequently, the Flan model receives a series of queries. The first checks for the presence of any biases in the text. If biases are detected, additional queries are deployed to locate, identify, and neutralize these biased elements. The results are then returned through the same path—first to the Lambda function, then through the API Gateway, and ultimately back to the user. If any bias was present in the original text, the user receives a comprehensive analysis indicating the types of biases detected, whether racial, ethnic, or gender. Specific sections of the text where these biases were found are highlighted, giving users a clear view of the changes made. Alongside this analysis, a new, de-biased version of their text is presented, effectively transforming potentially biased input into a more inclusive narrative.

In the following sections, we detail the steps to implement this solution.

Set up the React environment

We began by setting up our development environment for React. For bootstrapping a new React application with minimal configuration, we used create-react-app:

npx create-react-app my-app

Build the user interface

Using React, we designed a simple interface for users to input text, with a submission button, a reset button, and overlaying displays for presenting the processed results when they’re available.

Initiate the Flan model on SageMaker

We used SageMaker to create a pre-trained instance of the Flan language model with an endpoint for real-time inference. The model can be used against any JSON-structured payload like the following:

payload = {
      text_inputs: "text_inputs",
      max_length: <max_length>,
      num_return_sequences: <num_return_sequences>,
      top_k: <top_k>,
      top_p: <top_p>,
      do_sample: <do_sample>,
      num_beams: <num_beams>,
      seed: <seed>,
    };

Create a Lambda function

We developed a Lambda function that interacted directly with our SageMaker endpoint. The function was designed to receive a request with the user’s text, forward it to the SageMaker endpoint, and return the refined results, as shown in the following code (ENDPOINT_NAME was set up as the SageMaker instance endpoint):

import os
import io
import boto3
import json
import csv

# grab environment variables
ENDPOINT_NAME = os.environ['ENDPOINT_NAME']
runtime= boto3.client('runtime.sagemaker')

def lambda_handler(event, context):
    data = json.loads(json.dumps(event))
    payload = json.dumps(data['data']).encode('utf-8')

    query_response = runtime.invoke_endpoint(
        EndpointName=ENDPOINT_NAME,
        ContentType='application/json', 
        Body=payload)

    response_dict = json.loads(query_response['Body'].read())

    return response_dict['generated_texts']

Set up API Gateway

We configured a new REST API in API Gateway and linked it to our Lambda function. This connection allowed our React application to make HTTP requests to the API Gateway, which subsequently triggered the Lambda function.

Integrate the React app with the API

We updated the React application to make a POST request to the API Gateway when the submit button was clicked, with the body of the request being the user’s text. The JavaScript code we used to perform the API call is as follows (REACT_APP_AWS_ENDPOINT corresponds to the API Gateway endpoint bound to the Lambda call):

const makeAWSApiCall = (
    textInputs,
    maxLength,
    numReturnSequences,
    topK,
    topP,
    doSample,
    numBeams
  ) => {
    const axiosRequestUrl =
      `${process.env.REACT_APP_AWS_ENDPOINT}`;
    const requestData = {
      text_inputs: textInputs,
      max_length: maxLength,
      num_return_sequences: numReturnSequences,
      top_k: topK,
      top_p: topP,
      do_sample: doSample,
      num_beams: numBeams,
      seed: 8,
    };

    return axios.post(axiosRequestUrl, { data: requestData });
  };

Optimize prompt selection

To improve the accuracy of bias detection, we tested different prompts against the CrowS-Pairs dataset. Through this iterative process, we chose the prompts that gave us the highest accuracy.

Deploy and test the React app on Vercel

After building the application, we deployed it on Vercel to make it publicly accessible. We conducted extensive tests to ensure the application functioned as expected, from the user interface to the responses from the language model.

These steps laid the groundwork for creating our application for analyzing and de-biasing text. Despite the inherent complexity of the process, the use of tools like SageMaker, Lambda, and API Gateway streamlined the development, allowing us to focus on the core goal of the project—identifying and eliminating biases in text.

Conclusion

SageMaker JumpStart offers a convenient way to explore the features and capabilities of SageMaker. It provides curated one-step solutions, example notebooks, and deployable pre-trained models. These resources allow you to quickly learn and understand SageMaker. Additionally, you have the option to fine-tune the models and deploy them according to your specific needs. Access to JumpStart is available through Amazon SageMaker Studio or programmatically using the SageMaker APIs.

In this post, you learned how a student Hackathon team developed a solution in a short time using SageMaker JumpStart, which shows the potential of AWS and SageMaker JumpStart in enabling rapid development and deployment of sophisticated AI solutions, even by small teams or individuals.

To learn more about using SageMaker JumpStart, refer to Instruction fine-tuning for FLAN T5 XL with Amazon SageMaker Jumpstart and Zero-shot prompting for the Flan-T5 foundation model in Amazon SageMaker JumpStart.

ETH Analytics Club hosted ‘ETH Datathon,’ an AI/ML hackathon that draws more than 150 participants from ETH Zurich, University of Zurich, and EPFL. The event features workshops led by industry leaders, a 24-hour coding challenge, and valuable networking opportunities with fellow students and industry professionals. Great thanks to the ETH Hackathon team: Daniele Chiappalupi, Athina Nisioti, and Francesco Ignazio Re, as well as the rest of AWS organizing team: Alice Morano, Demir Catovic, Iana Peix, Jan Oliver Seidenfuss, Lars Nettemann, and Markus Winterholer.

The content and opinions in this post are those of the third-party author and AWS is not responsible for the content or accuracy of this post.

About the authors

Jun Zhang is a Solutions Architect based in Zurich. He helps Swiss customers architect cloud-based solutions to achieve their business potential. He has a passion for sustainability and strives to solve current sustainability challenges with technology. He is also a huge tennis fan and enjoys playing board games a lot.

Mohan Gowda leads Machine Learning team at AWS Switzerland. He works primarily with Automotive customers to develop innovative AI/ML solutions and platforms for next generation vehicles. Before working with AWS, Mohan worked with a Global Management Consulting firm with a focus on Strategy & Analytics. His passion lies in connected vehicles and autonomous driving.

Matthias Egli is the Head of Education in Switzerland. He is an enthusiastic Team Lead with a broad experience in business development, sales, and marketing.

Kemeng Zhang is an ML Engineer based in Zurich. She helps global customers design, develop, and scale ML-based applications to empower their digital capabilities to increase business revenue and reduce cost. She is also very passionate about creating human-centric applications by leveraging knowledge from behavioral science. She likes playing water sports and walking dogs.

Daniele Chiappalupi is a recent graduate from ETH Zürich. He enjoys every aspect of software engineering, from design to implementation, and from deployment to maintenance. He has a deep passion for AI and eagerly anticipates exploring, utilizing, and contributing to the latest advancements in the field. In his free time, he loves going snowboarding during colder months and playing pick-up basketball when the weather warms up.

Read More

We’re at an exciting inflection point in the widespread adoption of machine learning (ML), and we believe most customer experiences and applications will be reinvented with generative AI. Generative AI can create new content and ideas, including conversations, stories, images, videos, and music. Like most AI, generative AI is powered by ML models—very large models that are trained on vast amounts of data and commonly referred to as foundation models (FMs). FMs are based on transformers. Transformers are slow and memory-hungry on generating long text sequences due to the sheer size of the models. Large language models (LLMs) used to generate text sequences need immense amounts of computing power and have difficulty accessing the available high bandwidth memory (HBM) and compute capacity. This is because a large portion of the available memory bandwidth is consumed by loading the model’s parameters and by the auto-regressive decoding process.As a result, even with massive amounts of compute power, LLMs are limited by memory I/O and computation limits, preventing them from taking full advantage of the available hardware resources.

Overall, generative inference of LLMs has three main challenges (according to Pope et al. 2022):

• A large memory footprint due to massive model parameters and transient state during decoding. The parameters often exceed the memory of a single accelerator chip. Attention key-value caches also require substantial memory.
• Low parallelizability increases latency, especially with the large memory footprint, requiring substantial data transfers to load parameters and caches into compute cores each step. This results in high total memory bandwidth needs to meet latency targets.
• Quadratic scaling of attention mechanism compute relative to sequence length compounds the latency and computational challenges.

Batching is one of the techniques to address these challenges. Batching refers to the process of sending multiple input sequences together to a LLM and thereby optimizing the performance of the LLM inference. This approach helps improve throughput because model parameters don’t need to be loaded for every input sequence. The parameters can be loaded one time and used to process multiple input sequences. Batching efficiently utilizes the accelerator’s HBM bandwidth, resulting in higher compute utilization, improved throughput, and cost-effective inference.

This post examines techniques to maximize the throughput using batching techniques for parallelized generative inference in LLMs. We discuss different batching methods to reduce memory footprint, increase parallelizability, and mitigate the quadratic scaling of attention to boost throughput. The goal is to fully use hardware like HBM and accelerators to overcome bottlenecks in memory, I/O, and computation. Then we highlight how Amazon SageMaker large model inference (LMI) deep learning containers (DLCs) can help with these techniques. Finally, we present a comparative analysis of throughput improvements with each batching strategy on SageMaker using LMI DLCs to improve throughput for models like Llama v2. You can find an accompanying example notebook in the SageMaker examples GitHub repository.

Inferencing for large language models (LLMs)

Autoregressive decoding is the process by which language models like GPT generate text output one token at a time. It involves recursively feeding generated tokens back into the model as part of the input sequence in order to predict subsequent tokens. The steps are as follows:

1. The model receives the previous tokens in the sequence as input. For the first step, this is the starting prompt provided by the user.
2. The model predicts a distribution over the vocabulary for the next token.
3. The token with the highest predicted probability is selected and appended to the output sequence. Steps 2 and 3 are part of the decoding As of this writing, the most prominent decoding methods are greedy search, beam search, contrastive search, and sampling.
4. This new token is added to the input sequence for the next decoding step.
5. The model iterates through these steps, generating one new token per step, until an end-of-sequence marker is produced or the desired output length is reached.

Model serving for LLMs

Model serving for LLMs refers to the process of receiving input requests for text generation, making inferences, and returning the results to the requesting applications. The following are key concepts involved in model serving:

• Clients generate multiple inference requests, with each request consisting of sequence of tokens or input prompts
• Requests are received by the inference server (for example, DJLServing, TorchServe, Triton, or Hugging Face TGI)
• The inference server batches the inference requests and schedules the batch to the execution engine that includes model partitioning libraries (such as Transformers-NeuronX, DeepSpeed, Accelerate, or FasterTransformer) for running the forward pass (predicting the output token sequence) on the generative language model
• The execution engine generates response tokens and sends the response back to the inference server
• The inference server replies to the clients with the generated results

There are challenges with request-level scheduling when the inference server interacts with the execution engine at the request level, such as each request using a Python process, which requires a separate copy of model, which is memory restrictive. For example, as shown in the following figure, you can only accommodate to load a single copy of a model of size 80 GB on a machine learning (ML) instance with 96 GB of total accelerator device memory. You will need to load an additional copy of the entire model if you want to serve additional requests concurrently. This is not memory and cost efficient.

Now that we understand challenges posed by request-level scheduling, let’s look at different batching techniques that can help optimize throughput.

Batching techniques

In this section, we explain different batching techniques and show how to implement them using a SageMaker LMI container.

There are two main types of batching for inference requests:

• Client-side (static) – Typically, when a client sends a request to a server, the server will process each request sequentially by default, which is not optimal for throughput. To optimize the throughput, the client batches the inference requests in the single payload and the server implements the preprocessing logic to break down the batch into multiple requests and runs the inference for each request separately. In this option, the client needs to change the code for batching and the solution is tightly coupled with the batch size.
• Server-side (dynamic) – Another technique for batching is to use the inference to help achieve the batching on server side. As independent inference requests arrive at the server, the inference server can dynamically group them into larger batches on the server side. The inference server can manage the batching to meet a specified latency target, maximizing throughput while staying within the desired latency range. The inference server handles this automatically, so no client-side code changes are needed. The server-side batching includes different techniques to optimize the throughput further for generative language models based on the auto-regressive decoding. These batching techniques include dynamic batching, continuous batching, and PagedAttention (vLLM) batching.

Dynamic batching

Dynamic batching refers to combining the input requests and sending them together as a batch for inference. Dynamic batching is a generic server-side batching technique that works for all tasks, including computer vision (CV), natural language processing (NLP), and more.

In an LMI container, you can configure the batching of requests based on the following settings in serving.properties:

• batch_size – Refers to the size of the batch
• max_batch_delay – Refers to the maximum delay for batch aggregation

If either of these thresholds are met (meeting the maximum batch size or completion of the waiting period), then a new batch is prepared and pushed to the model for inferencing. The following diagram shows the dynamic batching of requests with different input sequence lengths being processed together by the model.

You can implement dynamic batching on SageMaker by configuring the LMI container’s serving.properties as follows:

#Dynamic Batching
engine=Python
option.entryPoint=djl_python.huggingface
batch_size=64 #example
max_batch_delay=1000 #example
option.tensor_parallel_degree=2 #example

Although dynamic batching can provide up to a four-times increase in throughput compared to no batching, we observe that GPU utilization is not optimal in this case because the system can’t accept another batch until all requests have completed processing.

Continuous batching

Continuous batching is an optimization specific for text generation. It improves throughput and doesn’t sacrifice the time to first byte latency. Continuous batching (also known as iterative or rolling batching) addresses the challenge of idle GPU time and builds on top of the dynamic batching approach further by continuously pushing newer requests in the batch. The following diagram shows continuous batching of requests. When requests 2 and 3 finish processing, another set of requests is scheduled.

The following interactive diagram dives deeper into how continuous batching works.

(Courtesy: https://github.com/InternLM/lmdeploy)

You can use a powerful technique to make LLMs and text generation efficient: caching some of the attention matrices. This means that the first pass of a prompt is different from the subsequent forward passes. For the first pass, you have to compute the entire attention matrix, whereas the follow-ups only require you to compute the new token attention. The first pass is called prefill throughout this code base, whereas the follow-ups are called decode. Because prefill is much more expensive than decode, we don’t want to do it all the time, but a currently running query is probably doing decode. If we want to use continuous batching as explained previously, we need to run prefill at some point in order to create the attention matrix required to be able to join the decode group.

This technique may allow up to a 20-times increase in throughput compared to no batching by effectively utilizing the idle GPUs.

You can fine-tune the following parameters in serving.properties of the LMI container for using continuous batching:

• engine – The runtime engine of the code. Values include Python, DeepSpeed, FasterTransformer, and MPI. Use MPI to enable continuous batching.
• rolling_batch – Enables iteration-level batching using one of the supported strategies. Values include auto, scheduler, and lmi-dist. We use lmi-dist for turning on continuous batching for Llama 2.
• max_rolling_batch_size – Limits the number of concurrent requests in the continuous batch. Defaults to 32.
• max_rolling_batch_prefill_tokens – Limits the number of tokens for caching. This needs to be tuned based on batch size and input sequence length to avoid GPU out of memory. It’s only supported for when rolling_batch=lmi-dist. Our recommendation is to set the value based on the number of concurrent requests x the memory required to store input tokens and output tokens per request.

The following is sample code for serving.properties for configuring continuous batching:

#Continuous Batching
engine=MPI
option.entryPoint=djl_python.huggingface
option.rolling_batch=auto
option.max_rolling_batch_size=64 #example
option.paged_attention=false
option.max_rolling_batch_prefill_tokens=16080 #example
option.tensor_parallel_degree=2 #example

PagedAttention batching

In the autoregressive decoding process, all the input tokens to the LLM produce their attention key and value tensors, and these tensors are kept in GPU memory to generate next tokens. These cached key and value tensors are often referred to as the KV cache or attention cache. As per the paper vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention, the KV cache takes up to 1.7 GB for a single sequence in Llama 13B. It is also dynamic. Its size depends on the sequence length, which is highly variable and unpredictable. As a result, efficiently managing the KV cache presents a significant challenge. The paper found that existing systems waste 60–80% of memory due to fragmentation and over-reservation.

PagedAttention is a new optimization algorithm developed by UC Berkeley that improves the continuous batching process by allowing the attention cache (KV cache) to be non-contiguous by allocating memory in fixed-size pages or blocks. This is inspired by virtual memory and paging concepts used by operating systems.

As per the vLLM paper, the attention cache of each sequence of tokens is partitioned into blocks and mapped to physical blocks through a block table. During the computation of attention, a PagedAttention kernel can use the block table to efficiently fetch the blocks from physical memory. This results in a significant reduction of memory waste and allows for larger batch size, increased GPU utilization, and higher throughput. The following figure illustrates partitioning the attention cache into non-contiguous pages.

The following diagram shows an inference example with PagedAttention. The key steps are:

1. The inference request is received with an input prompt.
2. In the prefill phase, attention is computed and key-values are stored in non-contiguous physical memory and mapped to logical key-value blocks. This mapping is stored in a block table.
3. The input prompt is run through the model (a forward pass) to generate the first response token. During the response token generation, the attention cache from the prefill phase is used.
4. During subsequent token generation, if the current physical block is full, additional memory is allocated in a non-contiguous fashion, allowing just-in-time allocation.

PagedAttention helps in near-optimal memory usage and reduction of memory waste. This allows for more requests to be batched together, resulting in a significant increase in throughput of inferencing.

The following code is a sample serving.properties for configuring PagedAttention batching in an LMI container on SageMaker:

#Paged Attention Batching
engine=MPI
option.entryPoint=djl_python.huggingface
option.rolling_batch=auto
option.max_rolling_batch_size=64 #example
option.paged_attention=true
option.max_rolling_batch_prefill_tokens=16080 #example
option.tensor_parallel_degree=2 #example

When to use which batching technique

The following figure summarizes the server-side batching techniques along with the sample serving.properties in LMI on SageMaker.

The following table summarizes the different batching techniques and their use cases.

Throughput comparison of different batching techniques for a large generative model on SageMaker

We performed performance benchmarking on a Llama v2 7B model on SageMaker using an LMI container and the different batching techniques discussed in this post with concurrent incoming requests of 50 and a total number of requests of 5,000.

We used three different input prompts of variable lengths for the performance test. In continuous and PagedAttention batching, the output tokens lengths were set to 64, 128, and 256 for the three input prompts, respectively. For dynamic batching, we used a consistent output token length of 128 tokens. We deployed SageMaker endpoints for the test with an instance type of ml.g5.24xlarge. The following table contains the results of the performance benchmarking tests.

We see an increase of approximately 2.3 times in throughput by using PagedAttention batching in comparison to dynamic batching for the Llama2-7B model on SageMaker using an LMI container.

Conclusion

In this post, we explained different batching techniques for LLMs inferencing and how it helps increase throughput. We showed how memory optimization techniques can increase the hardware efficiency by using continuous and PagedAttention batching and provide higher throughput values than dynamic batching. We saw an increase of approximately 2.3 times in throughput by using PagedAttention batching in comparison to dynamic batching for a Llama2-7B model on SageMaker using an LMI container. You can find the notebook used for testing the different batching techniques on GitHub.

About the authors

Gagan Singh is a Senior Technical Account Manager at AWS, where he partners with digital native startups to pave their path to heightened business success. With a niche in propelling Machine Learning initiatives, he leverages Amazon SageMaker, particularly emphasizing on Deep Learning and Generative AI solutions. In his free time, Gagan finds solace in trekking on the trails of the Himalayas and immersing himself in diverse music genres.

Dhawal Patel is a Principal Machine Learning Architect at AWS. He has worked with organizations ranging from large enterprises to mid-sized startups on problems related to distributed computing, and Artificial Intelligence. He focuses on Deep learning including NLP and Computer Vision domains. He helps customers achieve high performance model inference on SageMaker.

Venugopal Pai is a Solutions Architect at AWS. He lives in Bengaluru, India, and helps digital native customers scale and optimize their applications on AWS.

Read More

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.

Read More