Align Meta Llama 3 to human preferences with DPO, Amazon SageMaker Studio, and Amazon SageMaker Floor Fact

Massive language fashions (LLMs) have outstanding capabilities. Nonetheless, utilizing them in customer-facing functions usually requires tailoring their responses to align along with your group’s values and model identification. On this submit, we display the way to use direct desire optimization (DPO), a way that means that you can fine-tune an LLM with human desire information, along with Amazon SageMaker Studio and Amazon SageMaker Floor Fact to align the Meta Llama 3 8B Instruct mannequin responses to your group’s values.

Utilizing SageMaker Studio and SageMaker Floor Fact for DPO

With DPO, you may fine-tune an LLM with human desire information resembling rankings or rankings in order that it generates outputs that align to end-user expectations. DPO is computationally environment friendly and helps improve a mannequin’s helpfulness, honesty, and harmlessness, divert the LLM from addressing particular topics, and mitigate biases. On this method, you usually begin with deciding on an present or coaching a brand new supervised fine-tuned (SFT) mannequin. You employ the mannequin to generate responses and also you collect human suggestions on these responses. After that, you employ this suggestions to carry out DPO fine-tuning and align the mannequin to human preferences.

Whether or not you’re fine-tuning a pre-trained LLM with supervised fine-tuning (SFT) or loading an present fine-tuned mannequin for DPO, you usually want highly effective GPUs. The identical applies throughout DPO fine-tuning. With Amazon SageMaker, you may get began rapidly and experiment quickly through the use of managed Jupyter notebooks geared up with GPU situations. You may rapidly get began by making a JupyterLab house in SageMaker Studio, the built-in improvement surroundings (IDE) purpose-built for machine studying (ML), launch a JupyterLab software that runs on a GPU occasion.

Orchestrating the end-to-end information assortment workflow and creating an software for annotators to fee or rank mannequin responses for DPO fine-tuning will be time-consuming. SageMaker Floor Fact gives human-in-the-loop capabilities that assist you to arrange workflows, handle annotators, and acquire constant, high-quality suggestions.

This submit walks you thru the steps of utilizing DPO to align an SFT mannequin’s responses to the values of a fictional digital financial institution known as Instance Financial institution. Your pocket book runs in a JupyterLab house in SageMaker Studio powered by a single ml.g5.48xlarge occasion (8 A10G GPUs). Optionally, you may select to run this pocket book inside a smaller occasion kind resembling ml.g5.12xlarge (4 A10G GPUs) or ml.g6.12xlarge (4 L4 GPUs) with bitsandbytes quantization. You employ Meta Llama 3 8B Instruct (the Meta Llama 3 instruction tuned mannequin optimized for dialogue use instances from the Hugging Face Hub) to generate responses, SageMaker Floor Fact to gather desire information, and the DPOTrainer from the HuggingFace TRL library for DPO fine-tuning along with Parameter-Environment friendly High quality-Tuning (PEFT). You additionally deploy the aligned mannequin to a SageMaker endpoint for real-time inference. You should utilize the identical method with different fashions.

Answer overview

The next diagram illustrates the method.

The workflow accommodates the next key steps:

1. Load the Meta Llama 3 8B Instruct mannequin into SageMaker Studio and generate responses for a curated set of widespread and poisonous questions. The dataset serves because the preliminary benchmark for the mannequin’s efficiency.
2. The generated question-answer pairs are saved in Amazon Easy Storage Service (Amazon S3). These can be offered to the human annotators later to allow them to rank the mannequin responses.
3. Create a workflow in SageMaker Floor Fact to assemble human desire information for the responses. This includes creating a piece staff, designing a UI for suggestions assortment, and organising a labeling job.
4. Human annotators work together with the labeling portal to guage and rank the mannequin’s responses based mostly on their alignment to the group’s values.
5. The collected information is processed to stick to the DPOTrainer anticipated format.
6. Utilizing the Hugging Face TRL library and the DPOTrainer, fine-tune the Llama 3 mannequin utilizing the processed information from the earlier step.
7. Take a look at the fine-tuned mannequin on a holdout analysis dataset to evaluate its efficiency and confirm it meets the specified requirements.
8. Once you’re happy with the mannequin efficiency, you may deploy it to a SageMaker endpoint for real-time inference at scale.

Stipulations

To run the answer described on this submit, you could have an AWS account arrange, together with an AWS Id and Entry Administration (IAM) position that grants you the required permissions to create and entry the answer assets. In case you are new to AWS and haven’t created an account but, seek advice from Create a standalone AWS account.

To make use of SageMaker Studio, that you must have a SageMaker area arrange with a person profile that has the required permissions to launch the SageMaker Studio software. If you happen to’re new to SageMaker Studio, the Fast Studio setup is the quickest approach to get began. With a single click on, SageMaker provisions the required area with default presets, together with organising the person profile, IAM position, IAM authentication, and public web entry. The pocket book related to this submit assumes the usage of an ml.g5.48xlarge occasion kind. To assessment or enhance your quota limits, navigate to the AWS Service Quotas console, select AWS Companies within the navigation pane, select Amazon SageMaker, and seek advice from the worth for Studio JupyterLab Apps working on ml.g5.48xlarge situations.

Request a rise in quota worth larger than or equal to 1 for experimentation.

Meta Llama 3 8B Instruct is accessible underneath the Llama 3 license. To obtain the mannequin from Hugging Face, you want an entry token. If you happen to don’t have already got one, navigate to the Settings web page on the Hugging Face web site to acquire it.

Guarantee that the SageMaker Studio position has the required permissions for SageMaker Floor Fact and Amazon S3 entry. Once you’re working in SageMaker Studio, you’re already utilizing an IAM position, which you’ll want to change for launching SageMaker Floor Fact labeling jobs. To allow SageMaker Floor Fact performance, you need to connect the AWS managed coverage AmazonSageMakerGroundTruthExecution to your SageMaker Studio position. This coverage gives the important permissions for creating and managing labeling jobs.

For Amazon S3 entry, scoping permissions to particular buckets and actions enhances safety and aligns with greatest practices. This method adheres to the precept of least privilege, lowering potential dangers related to overly permissive insurance policies. The next is an instance of a restricted Amazon S3 coverage that grants solely the required permissions:

{
    "Model": "2012-10-17",
    "Assertion": [
        {
            "Effect": "Allow",
            "Action": [
                "s3:GetObject",
                "s3:PutObject",
                "s3:ListBucket"
            ],
            "Useful resource": [
                "arn:aws:s3:::<YOUR-BUCKET-NAME>",
                "arn:aws:s3:::<YOUR-BUCKET-NAME>/*"
            ]
        }
    ]
}

So as to add these insurance policies to your SageMaker Studio position, full the next steps:

1. On the IAM console, discover and select your SageMaker Studio position (it often begins with AmazonSageMaker-ExecutionRole-).
2. On the Permissions tab, select Add permissions after which Connect insurance policies.
3. Seek for and connect AmazonSageMakerGroundTruthExecution.
4. Create and connect the customized Amazon S3 inline coverage as proven within the previous instance, if wanted.

Keep in mind to observe the precept of least privilege, granting solely the permissions essential to your particular use case. Usually assessment your IAM roles and insurance policies to validate their alignment along with your safety necessities. For extra particulars on IAM insurance policies for SageMaker Floor Fact, seek advice from Use IAM Managed Insurance policies with Floor Fact.

Arrange the pocket book and surroundings

To get began, open SageMaker Studio and create a JupyterLab house. For Occasion, select ml.g5.48xlarge. Run the house, open JupyterLab, and clone the code within the following GitHub repository. You may configure the JupyterLab house to make use of as much as 100 GB in your Amazon Elastic Block Retailer (Amazon EBS) quantity. As well as, the ml.g5 occasion household comes with NVMe SSD native storage, which you should utilize within the JupyterLab software. The NVMe occasion retailer listing is mounted to the appliance container in /mnt/sagemaker-nvme. For this submit, you employ the NVMe storage out there within the ml.g5.48xlarge occasion.

When your house is prepared, clone the GitHub repo and open the pocket book llama3/rlhf-genai-studio/RLHF-with-Llama3-on-Studio-DPO.ipynb, which accommodates the answer code. Within the pop-up, make it possible for the Python 3 kernel is chosen.

Let’s undergo the pocket book. First, set up the required Python libraries:

import torch
import os
import sagemaker
import boto3
import datetime
from transformers import pipeline
import json
import asyncio
import aiofiles
from datasets import Dataset, load_dataset
from peft import (
get_peft_model,
    LoraConfig,
    prepare_model_for_kbit_training,
)
import bitsandbytes as bnb
from tqdm import tqdm
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    AutoModelForSequenceClassification
)
from IPython.core.show import show, HTML

The next line units the default path the place you retailer momentary artifacts to the placement within the NVMe storage:

cache_dir = "/mnt/sagemaker-nvme"

That is native storage, which implies that your information can be misplaced when the JupyterLab software is deleted, restarted, or patched. Alternatively, you may enhance your EBS quantity of your SageMaker Studio house to larger than or equal to 100 GB to supply adequate storage for the Meta Llama 3 base mannequin, PEFT adapter, and new merged fine-tuned mannequin.

Load Meta Llama 3 8B Instruct within the pocket book

After you may have imported the required libraries, you may obtain the Meta Llama 3 8B Instruct mannequin and its related tokenizers from Hugging Face:

base_model_id = "meta-llama/Meta-Llama-3-8B-Instruct"

mannequin = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    token=hf_access_token,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    cache_dir=cache_dir
)

mannequin.config.use_cache = False

tokenizer = AutoTokenizer.from_pretrained(
    base_model_id,
    token=hf_access_token,
    cache_dir=cache_dir
)

Gather preliminary mannequin responses for widespread and poisonous questions

The example_bank_questions.txt file accommodates a listing of widespread questions obtained by name facilities in monetary organizations mixed with a listing of poisonous and off-topic questions.

Earlier than you ask the mannequin to generate solutions to those questions, that you must specify the model and core values of Instance Financial institution. You’ll embody these values within the immediate as context later so the mannequin has the suitable info it wants to reply.

company_context = """Instance Financial institution is a next-generation digital financial institution on a mission to revolutionize the banking expertise. Based in 2020, we're dedicated to leveraging cutting-edge expertise to make banking easy, accessible, and clear for everybody. In Instance Financial institution, we consider that banking needs to be seamless, intuitive, and tailor-made to the wants of contemporary shoppers. Our founders, seasoned professionals from the tech and finance industries, got down to create a financial institution that places folks first, empowering them to take management of their funds with ease. At Instance Financial institution, we envision a world the place banking is now not a chore however a pleasant expertise. We're devoted to breaking down limitations and democratizing entry to monetary providers. Our aim is to empower people and companies alike by offering them with the instruments and assets they should thrive in an more and more digital panorama.
Our values:
- Innovation: We embrace cutting-edge applied sciences and repeatedly search out modern options to ship the absolute best banking expertise. We're a digital-only financial institution, which implies we haven't any bodily branches. As a substitute, we provide all of our providers on-line or by way of our cellular app. This enables us to maintain our prices low and go the financial savings on to our prospects.
- Transparency: We're dedicated to being direct and trustworthy with our prospects. We consider that transparency is vital to constructing belief, and we would like our prospects to really feel assured that they're making knowledgeable selections about their cash. That is why we offer clear and concise details about our services and products, and we're at all times out there to reply any questions our prospects could have.
- Accessibility: Our providers are designed to be inclusive and user-friendly, catering to a various vary of shoppers, no matter their monetary backgrounds.
- Safety: We prioritize the security and safety of our prospects' information and property, using state-of-the-art encryption and cybersecurity measures.
Along with our core values, Instance Financial institution gives a variety of modern monetary services and products:
- Loans: Whether or not you’re trying to purchase a house, begin a enterprise, or finance a serious buy, our versatile mortgage choices are designed to fulfill your wants. With aggressive rates of interest and a easy software course of, acquiring a mortgage has by no means been simpler.
- Credit score Playing cards: Our bank cards include a bunch of advantages together with cashback rewards, low-interest charges, and no annual charges. Handle your spending effortlessly with real-time notifications and intuitive budgeting instruments.
- Cellular Apps: Our user-friendly apps on the Google Play Retailer and Apple App Retailer supply a seamless banking expertise. From checking balances to transferring funds, our apps guarantee you may have full management of your funds at your fingertips.
- Financial savings and Investments: Develop your wealth with our high-yield financial savings accounts and a wide range of funding choices. Our monetary advisors can be found that will help you make knowledgeable selections tailor-made to your monetary objectives.
- Buyer Help: We offer 24/7 buyer help to help with any inquiries or points. Our devoted staff is at all times prepared to assist, making certain you obtain the absolute best service always.
At Instance Financial institution, we're dedicated to enhancing your monetary well-being by way of innovation, transparency, and unparalleled service. Be part of us at the moment and expertise the way forward for banking.
"""

Now you’re able to invoke the mannequin. For every query within the file, you assemble a immediate that accommodates the context and the precise query. You ship the immediate to the mannequin 4 occasions to generate 4 totally different outputs and save the ends in the llm_responses.json file.

questions="example_bank_questions.txt"
llm_responses = os.path.be part of(sample_files_path, 'llm_responses.json')

from timeit import default_timer as timer
import tqdm.asyncio

async def invoke_model(query, context):
    pipe = pipeline("text-generation", mannequin=mannequin, tokenizer=tokenizer)
    messages = [
        {"role": "user", "content": f"{context}: {question}"}
    ]

    terminators = [
        tokenizer.eos_token_id,
        tokenizer.convert_tokens_to_ids("<|eot_id|>")
    ]

    response = pipe(
        messages, 
        max_new_tokens=120, 
        do_sample=True,
        temperature=gl_temperature, 
        top_p=gl_top_p, 
        eos_token_id=terminators
    )[0]['generated_text'][-1]
    return response['content']

async def process_lines(file_path):
    outcomes = []
    context = f"""{company_context} You're a customer support agent for {company_name} Generally you're sensible along with your solutions. Reply the next buyer query in a single or two sentences:
    """
    async with aiofiles.open(file_path, 'r') as file:
        strains = [line async for line in file]
        for line in tqdm.asyncio.tqdm(strains, desc="Processing Query Financial institution"):
            begin = timer()
            responses = await asyncio.collect(*[invoke_model(line, context) for _ in range(4)])
            consequence = {
                'context': context,
                'query': line.strip(),
                'responses': responses
            }
            finish = timer()
            outcomes.append(consequence)
    return outcomes

outcomes = await process_lines(questions)

with open(llm_responses, 'w') as file:
    json.dump(
        outcomes, 
        file, 
        indent=4
    )

The next is an instance entry from llm_reponses.json.

Arrange the SageMaker Floor Fact labeling job and human desire information

To fine-tune the mannequin utilizing DPO, that you must collect human desire information for the generated responses. SageMaker Floor Fact helps orchestrate the information assortment course of. It gives customizable labeling workflows and sturdy workforce administration options for rating duties. This part reveals you the way to arrange a SageMaker Floor Fact labeling job and invite a human workforce with requisite experience to assessment the LLM responses and rank them.

Arrange the workforce

A non-public workforce in SageMaker Floor Fact consists of people who’re particularly invited to carry out information labeling duties. These people will be workers or contractors who’ve the required experience to guage the mannequin’s responses. Organising a non-public workforce helps obtain information safety and high quality by limiting entry to trusted people for information labeling.

For this use case, the workforce consists of the group of people that will rank the mannequin responses. You may arrange a non-public workforce utilizing the SageMaker console by creating a non-public staff and alluring members by way of electronic mail. For detailed directions, seek advice from Create a Non-public Workforce (Amazon SageMaker Console).

Create the instruction template

With the instruction template, you may handle the UI and information human annotators in reviewing mannequin outputs. It wants to obviously current the mannequin responses and supply an easy method for the annotators to rank them. Right here, you employ the textual content rating template. This template means that you can show the directions for the human reviewer and the prompts with the pregenerated LLM responses. The annotator evaluations the immediate and responses and ranks the latter based mostly on their alignment to the group’s model.

The definition of the template is as follows. The template reveals a pane on the left with directions from the job requester, a immediate on the high, and three LLM responses in the principle physique. The precise aspect of the UI is the place the annotator ranks the responses from most to least preferable.

  <html>
  <head>
    <meta charset="UTF-8" />
    <hyperlink rel="stylesheet" href="https://property.crowd.aws/css/gen-ai-components.css" />
    <hyperlink rel="icon" href="information:picture/svg+xml,<svg xmlns=%22http://www.w3.org/2000/svgpercent22 viewBox=%220 0 100 100percent22><textual content y=%22.9empercent22 font-size=%2290percent22>🥇</textual content></svg>" />
    <title>Textual content Rating Software</title>
    <script src="https://property.crowd.aws/gen-ai-components.js"></script>
  </head>

  <physique>
    <div>
      <crowd-text-ranking
        crowd-form-element-id="crowd-form-submit"
        directions="Rank the next responses from a language mannequin in keeping with their alignment to the organisation"s model.'
        ordinal-ranking-dimensions="[{"name":"BrandValue","allowTie":true}]"
        textual content="{{ job.enter.supply }}"
        responses="{ to_json }" />
    </div>
    <crowd-form id="crowd-form-submit" fashion="show: none"></crowd-form>
    <script src="https://property.crowd.aws/crowd-html-elements.js"></script>
  </physique>
</html>

The template is saved regionally in your Studio JupyterLab house EBS quantity as directions.template in a short lived listing. You then add this template file to your designated S3 bucket utilizing s3.upload_file(), putting it within the specified bucket and prefix. This Amazon S3 hosted template can be referenced once you create the SageMaker Floor Fact labeling job, so employees see the proper interface for the textual content rating job.

Preprocess the enter information

Earlier than you create the labeling job, confirm that the enter information matches the format anticipated by SageMaker Floor Fact and is saved as a JSON file in Amazon S3. You should utilize the prompts and responses within the llm_responses.json file to create the manifest file inp-manifest-trank.json. Every row within the manifest file accommodates a JSON object (source-responses pair). The earlier entry now seems to be like the next code.

Add the structured information to the S3 bucket in order that it may be ingested by SageMaker Floor Fact.

Create the labeling job

Now you’re able to configure and launch the labeling job utilizing the SageMaker API from throughout the pocket book. This includes specifying the work staff, UI template, and information saved within the S3 bucket. By setting applicable parameters resembling job deadlines and the variety of employees per information object, you may run jobs effectively and successfully. The next code reveals the way to begin the labeling job:

sm_client.create_labeling_job(
    LabelingJobName=labeling_job_name,
    LabelAttributeName="label",
    InputConfig={
        'DataSource': {
            'S3DataSource': {
                'ManifestS3Uri': model_responses_s3_uri
            }
        }
    },
    OutputConfig={
        'S3OutputPath': 's3://{}/{}/output/'.format(bucket,prefix) #Enter S3 URI of Output folder
    },
    RoleArn=position, 
    HumanTaskConfig={
        'WorkteamArn': WORKTEAM_ARN,
        'UiConfig':{
            'UiTemplateS3Uri': UI_TEMPLATE_S3_URI
        },
        'PreHumanTaskLambdaArn': 'arn:aws:lambda:us-east-1:432418664414:operate:PRE-PassThrough',
        'TaskKeywords': [
            'QnA',
        ],
        'TaskTitle': 'Rank LLM responses',
        'TaskDescription': "Rank the responses offered by the LLM",
        'NumberOfHumanWorkersPerDataObject': 1,
        'TaskTimeLimitInSeconds': 60*30,
        'TaskAvailabilityLifetimeInSeconds': 60*60*24*10,
        'MaxConcurrentTaskCount': 100,
        'AnnotationConsolidationConfig': {
            'AnnotationConsolidationLambdaArn': 'arn:aws:lambda:us-east-1:432418664414:operate:ACS-PassThrough'
        } 
    }

Because the job is launched, it’s important to watch its progress carefully, ensuring duties are being distributed and accomplished as anticipated.

Collect human suggestions by way of the labeling portal

When the job setup is full, annotators can log in to the labeling portal and begin rating the mannequin responses.

Employees can first seek the advice of the Directions pane to know the duty, then use the principle interface to guage and rank the mannequin’s responses in keeping with the given standards. The next screenshot illustrates the UI.

The human suggestions is collected and saved in an S3 bucket. This suggestions would be the foundation for DPO. With this information, you’ll fine-tune the Meta Llama 3 mannequin and align its responses with the group’s values, bettering its general efficiency.

Align Meta Llama 3 8B Instruct with the DPOTrainer

On this part, we present the way to use the desire dataset that you simply ready utilizing SageMaker Floor Fact to fine-tune the mannequin utilizing DPO. DPO explicitly optimizes the mannequin’s output based mostly on human evaluations. It aligns the mannequin’s conduct extra carefully with human expectations and improves its efficiency on duties requiring nuanced understanding and contextual appropriateness. By integrating human preferences, DPO enhances the mannequin’s relevance, coherence, and general effectiveness in producing desired responses.

DPO makes it extra easy to preference-tune a mannequin compared to different standard strategies resembling Proximal Coverage Optimization (PPO). DPO eliminates the need for a separate rewards mannequin, thereby avoiding the associated fee related to coaching it. Moreover, DPO requires considerably much less information to attain efficiency akin to PPO.

High quality-tuning a language mannequin utilizing DPO consists of two steps:

1. Collect a desire dataset with optimistic and adverse chosen pairs of technology, given a immediate.
2. Maximize the log-likelihood of the DPO loss immediately.

To study extra in regards to the DPO algorithm, seek advice from the next whitepaper.

Anticipated information format

The DPO coach expects a really particular format for the dataset, which accommodates sentence pairs the place one sentence is a selected response and the opposite is a rejected response. That is represented as a Python dictionary with three keys:

• immediate – Consists of the context immediate given to a mannequin at inference time for textual content technology
• chosen – Comprises the popular generated response to the corresponding immediate
• rejected – Comprises the response that’s not most popular or shouldn’t be the sampled response for the given immediate

The next operate definition illustrates the way to course of the information saved in Amazon S3 to create a DPO dataset utilizing with pattern pairs and a immediate:

def return_prompt_and_responses(samples, index):
    immediate = f"{samples['context']}nn{samples['question']}"
    chosen_index = response_rankings[index]["responseRankings"].index(1)
    rejected_index = response_rankings[index]["responseRankings"].index(4)

    immediate = {"position": "person", "content material": immediate},

    chosen_messages = [
        {"role": "assistant", "content": samples["responses"][chosen_index]},
    ]
    rejected_messages = [
        # {"role": "system", "content": prompt},
        {"role": "assistant", "content": samples["responses"][rejected_index]},
    ]
    
    return {
        "immediate": tokenizer.apply_chat_template(immediate, tokenize=False),
        "chosen": "{}".format(tokenizer.apply_chat_template(chosen_messages, tokenize=False).substitute('<|begin_of_text|>', '')),
        "rejected": "{}".format(tokenizer.apply_chat_template(rejected_messages, tokenize=False).substitute('<|begin_of_text|>', ''))
    }

Right here is an instance sentence pair:

You cut up the DPO coach dataset into prepare and check samples utilizing an 80/20 cut up and tokenize the dataset in preparation for DPO fine-tuning:

dataset = prepared_dataset.train_test_split(test_size=0.2)

dataset["train"].to_json(
    os.path.be part of(sample_files_path, "processed_human_feedback", "train_dataset.json"), 
    orient="data", 
    index="False"
)

dataset["test"].to_json(
    os.path.be part of(sample_files_path, "processed_human_feedback", "test_dataset.json"), 
    orient="data", 
    index="False"

Supervised fine-tuning utilizing DPO

Now that the dataset is formatted for the DPO coach, you should utilize the prepare and check datasets ready earlier to provoke the DPO mannequin fine-tuning. Meta Llama 3 8B belongs to a class of small language fashions, however even Meta Llama 3 8B barely matches right into a SageMaker ML occasion like ml.g5.48xlarge in fp16 or fp32, leaving little room for full fine-tuning. You should utilize PEFT with DPO to fine-tune Meta Llama 3 8B’s responses based mostly on human preferences. PEFT is a technique of fine-tuning that focuses on coaching solely a subset of the pre-trained mannequin’s parameters. This method includes figuring out a very powerful parameters for the brand new job and updating solely these parameters throughout coaching. By doing so, PEFT can considerably cut back the computation required for fine-tuning. See the next code:

# configure PEFT module
peft_config = LoraConfig(
    r=512,
    lora_alpha=1024,
    lora_dropout=0.05,
    task_type="CAUSAL_LM",
    target_modules="all-linear",

For a full checklist of LoraConfig coaching arguments, seek advice from LoRA. At a excessive degree, that you must initialize the DPOTrainer with the next parts: the mannequin you wish to prepare, a reference mannequin (ref_model) used to calculate the implicit rewards of the popular and rejected responses, the beta hyperparameter that controls the steadiness between the implicit rewards assigned to the popular and rejected responses, and a dataset containing immediate, chosen, and rejected responses. If ref_model=None, the coach will create a reference mannequin with the identical structure because the enter mannequin to be optimized. See the next code:

from trl import DPOConfig, DPOTrainer

dpo_model_dir = "/path/to/save/dpo/mannequin"

args = DPOConfig(
    output_dir=dpo_model_dir,               # listing to avoid wasting and repository id
    num_train_epochs=5,                     # variety of coaching epochs
    per_device_train_batch_size=2,
    gradient_accumulation_steps=1,
    gradient_checkpointing=True,            # use gradient checkpointing to avoid wasting reminiscence
    optim = "adamw_torch_fused",            # use fused adamw optimizer
    learning_rate=1e-5,                     # 10x increased LR than QLoRA paper
    max_grad_norm=0.3,                      # max gradient norm based mostly on QLoRA paper
    warmup_ratio=0.1,                       # warmup ratio based mostly on QLoRA paper
    lr_scheduler_type="cosine",             # use cosine studying fee scheduler
    logging_steps=10,                       
    save_steps=10,                         # when to avoid wasting checkpoint
    evaluation_strategy="steps",            
    eval_steps=100,
    bf16=True,                              # use bfloat16 precision
    tf32=True,                              # use tf32 precision
    push_to_hub=False,                      # push mannequin to hub,
    report_to='tensorboard',
    remove_unused_columns=False
)

dpo_args = {
    "beta": 0.1,                            # The beta consider DPO loss. Increased beta means much less divergence
    "loss_type": "sigmoid"                  # The loss kind for DPO.
}

coach = DPOTrainer(
    mannequin,
    ref_model=None,
    peft_config=peft_config,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
    max_length=max_seq_length,
    max_prompt_length=prompt_length,
    beta=dpo_args["beta"],
    loss_type=dpo_args["loss_type"],
)

# kick-off mannequin coaching
coach.prepare()

When you begin the coaching, you may see the standing within the pocket book:

When mannequin fine-tuning is full, save the PEFT adapter mannequin to disk and merge it with the bottom mannequin to create a newly tuned mannequin. You should utilize the saved mannequin for native inference and validation or deploy it as a SageMaker endpoint after you may have gained adequate confidence within the mannequin’s responses.

peft_output_dir = "/path/to/save/tuned/mannequin/"
print(f"saving peft mannequin to: {peft_output_dir}")
coach.save_model(output_dir=peft_output_dir)
...
...
merged_model = mannequin.merge_and_unload()
...
...
merged_model.save_pretrained(
    new_dpo_output_dir,
    safe_serialization=True,
    max_shard_size="9GB"
)

Consider the fine-tuned mannequin inside a SageMaker Studio pocket book

Earlier than you host your mannequin for inference, confirm that its response optimization aligns with person preferences. You may acquire the mannequin’s response each earlier than and after DPO fine-tuning and examine them aspect by aspect, as proven within the following desk.

The DPO Mannequin Response column signifies the RLHF aligned mannequin’s response post-fine-tuning, and the Rejected Mannequin Response column refers back to the mannequin’s response to the enter immediate previous to fine-tuning.

Deploy the mannequin to a SageMaker endpoint

After you may have gained adequate confidence in your mannequin, you may deploy it to a SageMaker endpoint for real-time inference. SageMaker endpoints are totally managed and supply auto scaling capabilities. For this submit, we use DJL Serving to host the fine-tuned, DPO-aligned Meta Llama3 8B mannequin. To study extra about internet hosting your LLM utilizing DJL Serving, seek advice from Deploy giant fashions on Amazon SageMaker utilizing DJLServing and DeepSpeed mannequin parallel inference.

To deploy an LLM immediately out of your SageMaker Studio pocket book utilizing DJL Serving, full the next steps:

1. Add mannequin weights and different mannequin artifacts to Amazon S3.
2. Create a meta-model definition file known as serving.properties. This definition file dictates how the DJL Serving container is configured for inference.

engine = DeepSpeed
possibility.tensor_parallel_degree = 1
possibility.s3url = s3://<MY-TEST-BUCKET>/llama3-dpo-ft/modelweights
possibility.hf_access_token=hf_xx1234

3. Create a customized inference file known as mannequin.py, which defines a customized inference logic:

%%writefile llama3-serving-model/mannequin.py

from djl_python import Enter, Output
...

predictor = None


def get_model(properties):

    ...
    return generator


def deal with(inputs: Enter) -> None:
    ...
    outputs = predictor(message, **generation_kwargs)[0]['generated_text'][-1]
    consequence = {"outputs": outputs['content']}
    return Output().add(consequence)

4. Deploy the DPO fine-tuned mannequin as a SageMaker endpoint:

from sagemaker import image_uris
from sagemaker.mannequin import Mannequin
from datetime import datetime

inference_image_uri = image_uris.retrieve(
    framework="djl-deepspeed",
    area=area,
    model="0.23.0"
)

...

dpo_model.deploy(
    initial_instance_count=1,
    instance_type="ml.g5.2xlarge",
    endpoint_name=f"ep-{dpo_model.title}",
    container_startup_health_check_timeout=900,
    wait=False, # <-- Set to True, when you would favor to attend 6-8 minutes for the endpoint to spin up
)

5. Invoke the hosted mannequin for inference utilizing the sageMaker.Predictor class:

dpo_ft_predictor = sagemaker.Predictor(
    endpoint_name="my_custom_dpo_endpoint",
    sagemaker_session=sess,
    serializer=serializers.JSONSerializer(),
    deserializer=deserializers.JSONDeserializer(),
)
...
# invoke inference
response = dpo_ft_predictor.predict(
    {
        "inputs": content material,
        "parameters": parameters
    }
)

Clear up

After you full your duties within the SageMaker Studio pocket book, bear in mind to cease your JupyterLab workspace to forestall incurring further costs. You are able to do this by selecting Cease subsequent to your JupyterLab house. Moreover, you may have the choice to arrange lifecycle configuration scripts that may routinely shut down assets after they’re not in use.

If you happen to deployed the mannequin to a SageMaker endpoint, run the next code on the finish of the pocket book to delete the endpoint:

#delete your endpoint
sm_client.delete_endpoint(EndpointName=tg_sm_model.endpoint_name)

Conclusion

Amazon SageMaker gives instruments to streamline the method of fine-tuning LLMs to align with human preferences. With SageMaker Studio, you may experiment interactively with totally different fashions, questions, and fine-tuning strategies. With SageMaker Floor Fact, you may arrange workflows, handle groups, and acquire constant, high-quality human suggestions.

On this submit, we confirmed the way to improve the efficiency of Meta Llama 3 8B Instruct by fine-tuning it utilizing DPO on information collected with SageMaker Floor Fact. To get began, launch SageMaker Studio and run the pocket book out there within the following GitHub repo. Share your ideas within the feedback part!

In regards to the Authors

Anastasia Tzeveleka is a GenAI/ML Specialist Options Architect at AWS. As a part of her work, she helps prospects construct basis fashions and create scalable generative AI and machine studying options utilizing AWS providers.

Pranav Murthy is an AI/ML Specialist Options Architect at AWS. He focuses on serving to prospects construct, prepare, deploy and migrate machine studying (ML) workloads to SageMaker. He beforehand labored within the semiconductor trade creating giant pc imaginative and prescient (CV) and pure language processing (NLP) fashions to enhance semiconductor processes. In his free time, he enjoys enjoying chess and touring.

Sundar Raghavan is an AI/ML Specialist Options Architect at AWS, serving to prospects construct scalable and cost-efficient AI/ML pipelines with Human within the Loop providers. In his free time, Sundar loves touring, sports activities and having fun with outside actions together with his household.