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

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

Utilizing SageMaker Studio and SageMaker Floor Reality for DPO

With DPO, you may fine-tune an LLM with human desire knowledge comparable to scores 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 sometimes begin with choosing an present or coaching a brand new supervised fine-tuned (SFT) mannequin. You utilize 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 might be fine-tuning a pre-trained LLM with supervised fine-tuning (SFT) or loading an present fine-tuned mannequin for DPO, you sometimes want highly effective GPUs. The identical applies throughout DPO fine-tuning. With Amazon SageMaker, you will get began shortly and experiment quickly through the use of managed Jupyter notebooks outfitted with GPU cases. You may shortly get began by making a JupyterLab house in SageMaker Studio, the built-in growth atmosphere (IDE) purpose-built for machine studying (ML), launch a JupyterLab software that runs on a GPU occasion.

Orchestrating the end-to-end knowledge assortment workflow and growing an software for annotators to fee or rank mannequin responses for DPO fine-tuning will be time-consuming. SageMaker Floor Reality presents human-in-the-loop capabilities that enable you arrange workflows, handle annotators, and gather 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 referred to 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 sort comparable to ml.g5.12xlarge (4 A10G GPUs) or ml.g6.12xlarge (4 L4 GPUs) with bitsandbytes quantization. You utilize Meta Llama 3 8B Instruct (the Meta Llama 3 instruction tuned mannequin optimized for dialogue use circumstances from the Hugging Face Hub) to generate responses, SageMaker Floor Reality to gather desire knowledge, and the DPOTrainer from the HuggingFace TRL library for DPO fine-tuning along with Parameter-Environment friendly Tremendous-Tuning (PEFT). You additionally deploy the aligned mannequin to a SageMaker endpoint for real-time inference. You should use the identical strategy with different fashions.

Resolution overview

The next diagram illustrates the strategy.

The workflow incorporates the next key steps:

1. Load the Meta Llama 3 8B Instruct mannequin into SageMaker Studio and generate responses for a curated set of frequent 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 shall be offered to the human annotators later to allow them to rank the mannequin responses.
3. Create a workflow in SageMaker Floor Reality to collect human desire knowledge for the responses. This entails creating a piece workforce, designing a UI for suggestions assortment, and establishing a labeling job.
4. Human annotators work together with the labeling portal to judge and rank the mannequin’s responses primarily based on their alignment to the group’s values.
5. The collected knowledge 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 knowledge from the earlier step.
7. Check the fine-tuned mannequin on a holdout analysis dataset to evaluate its efficiency and confirm it meets the specified requirements.
8. While you’re glad with the mannequin efficiency, you may deploy it to a SageMaker endpoint for real-time inference at scale.

Conditions

To run the answer described on this submit, you have to have an AWS account arrange, together with an AWS Id and Entry Administration (IAM) function that grants you the mandatory permissions to create and entry the answer sources. 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 mandatory permissions to launch the SageMaker Studio software. In case you’re new to SageMaker Studio, the Fast Studio setup is the quickest solution to get began. With a single click on, SageMaker provisions the required area with default presets, together with establishing the person profile, IAM function, IAM authentication, and public web entry. The pocket book related to this submit assumes the usage of an ml.g5.48xlarge occasion sort. To overview or improve 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 cases.

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

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

Be sure that the SageMaker Studio function has the mandatory permissions for SageMaker Floor Reality and Amazon S3 entry. While you’re working in SageMaker Studio, you’re already utilizing an IAM function, which you’ll want to change for launching SageMaker Floor Reality labeling jobs. To allow SageMaker Floor Reality performance, it’s best to connect the AWS managed coverage AmazonSageMakerGroundTruthExecution to your SageMaker Studio function. This coverage supplies 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 strategy adheres to the precept of least privilege, decreasing potential dangers related to overly permissive insurance policies. The next is an instance of a restricted Amazon S3 coverage that grants solely the mandatory permissions:

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

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

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

Keep in mind to comply with the precept of least privilege, granting solely the permissions obligatory on your particular use case. Commonly overview your IAM roles and insurance policies to validate their alignment together with your safety necessities. For extra particulars on IAM insurance policies for SageMaker Floor Reality, seek advice from Use IAM Managed Insurance policies with Floor Reality.

Arrange the pocket book and atmosphere

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 applying container in /mnt/sagemaker-nvme. For this submit, you employ the NVMe storage obtainable 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 incorporates the answer code. Within the pop-up, be sure that the Python 3 kernel is chosen.

Let’s undergo the pocket book. First, set up the mandatory 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 short-term artifacts to the situation within the NVMe storage:

cache_dir = "/mnt/sagemaker-nvme"

That is native storage, which signifies that your knowledge shall be misplaced when the JupyterLab software is deleted, restarted, or patched. Alternatively, you may improve your EBS quantity of your SageMaker Studio house to better than or equal to 100 GB to supply enough 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’ve got imported the mandatory 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
)

Accumulate preliminary mannequin responses for frequent and poisonous questions

The example_bank_questions.txt file incorporates a listing of frequent questions acquired 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 know-how to make banking easy, accessible, and clear for everybody. In Instance Financial institution, we consider that banking must 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 companies. Our aim is to empower people and companies alike by offering them with the instruments and sources they should thrive in an more and more digital panorama.
Our values:
- Innovation: We embrace cutting-edge applied sciences and repeatedly hunt down progressive options to ship the absolute best banking expertise. We're a digital-only financial institution, which suggests we haven't any bodily branches. As a substitute, we provide all of our companies on-line or via our cell app. This permits us to maintain our prices low and move the financial savings on to our prospects.
- Transparency: We're dedicated to being direct and sincere with our prospects. We consider that transparency is vital to constructing belief, and we wish our prospects to really feel assured that they're making knowledgeable choices about their cash. That is why we offer clear and concise details about our services and products, and we're at all times obtainable to reply any questions our prospects might have.
- Accessibility: Our companies are designed to be inclusive and user-friendly, catering to a various vary of shoppers, no matter their monetary backgrounds.
- Safety: We prioritize the protection and safety of our prospects' knowledge and belongings, using state-of-the-art encryption and cybersecurity measures.
Along with our core values, Instance Financial institution presents a spread of progressive monetary services and products:
- Loans: Whether or not you’re seeking to purchase a house, begin a enterprise, or finance a significant 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 number 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've got full management of your funds at your fingertips.
- Financial savings and Investments: Develop your wealth with our high-yield financial savings accounts and quite a lot of funding choices. Our monetary advisors can be found that will help you make knowledgeable choices tailor-made to your monetary objectives.
- Buyer Help: We offer 24/7 buyer assist to help with any inquiries or points. Our devoted workforce 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 via innovation, transparency, and unparalleled service. Be a part of us in the present day 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 incorporates the context and the precise query. You ship the immediate to the mannequin 4 occasions to generate 4 totally different outputs and save the leads to 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} Typically you might be good together 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)])
            outcome = {
                'context': context,
                'query': line.strip(),
                'responses': responses
            }
            finish = timer()
            outcomes.append(outcome)
    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 Reality labeling job and human desire knowledge

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

Arrange the workforce

A non-public workforce in SageMaker Floor Reality consists of people who’re particularly invited to carry out knowledge labeling duties. These people will be workers or contractors who’ve the required experience to judge the mannequin’s responses. Organising a non-public workforce helps obtain knowledge safety and high quality by limiting entry to trusted people for knowledge 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 workforce and alluring members via e 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 manner for the annotators to rank them. Right here, you employ the textual content rating template. This template permits you to 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 primarily based on their alignment to the group’s model.

The definition of the template is as follows. The template exhibits 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 appropriate aspect of the UI is the place the annotator ranks the responses from most to least preferable.

The template is saved domestically in your Studio JupyterLab house EBS quantity as directions.template in a brief listing. Then you definitely add this template file to your designated S3 bucket utilizing s3.upload_file(), inserting it within the specified bucket and prefix. This Amazon S3 hosted template shall be referenced if you create the SageMaker Floor Reality labeling job, so employees see the right interface for the textual content rating job.

Preprocess the enter knowledge

Earlier than you create the labeling job, confirm that the enter knowledge matches the format anticipated by SageMaker Floor Reality and is saved as a JSON file in Amazon S3. You should use 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 incorporates a JSON object (source-responses pair). The earlier entry now seems like the next code.

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

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 entails specifying the work workforce, UI template, and knowledge saved within the S3 bucket. By setting applicable parameters comparable to job closing dates and the variety of employees per knowledge object, you may run jobs effectively and successfully. The next code exhibits the right 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=function, 
    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 supplied 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 observe its progress intently, ensuring duties are being distributed and accomplished as anticipated.

Collect human suggestions via 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 grasp the duty, then use the principle interface to judge and rank the mannequin’s responses based on 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 knowledge, you’ll fine-tune the Meta Llama 3 mannequin and align its responses with the group’s values, bettering its total efficiency.

Align Meta Llama 3 8B Instruct with the DPOTrainer

On this part, we present the right way to use the desire dataset that you simply ready utilizing SageMaker Floor Reality to fine-tune the mannequin utilizing DPO. DPO explicitly optimizes the mannequin’s output primarily based on human evaluations. It aligns the mannequin’s conduct extra intently 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 total effectiveness in producing desired responses.

DPO makes it extra easy to preference-tune a mannequin compared to different widespread methods comparable to Proximal Coverage Optimization (PPO). DPO eliminates the need for a separate rewards mannequin, thereby avoiding the price related to coaching it. Moreover, DPO requires considerably much less knowledge to realize efficiency akin to PPO.

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

1. Collect a desire dataset with optimistic and destructive chosen pairs of era, given a immediate.
2. Maximize the log-likelihood of the DPO loss straight.

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

Anticipated knowledge format

The DPO coach expects a really particular format for the dataset, which incorporates 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 era
• chosen – Incorporates the popular generated response to the corresponding immediate
• rejected – Incorporates the response that isn’t most well-liked or shouldn’t be the sampled response for the given immediate

The next operate definition illustrates the right 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 = {"function": "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 practice 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="information", 
    index="False"
)

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

Supervised fine-tuning utilizing DPO

Now that the dataset is formatted for the DPO coach, you should utilize the practice 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 use PEFT with DPO to fine-tune Meta Llama 3 8B’s responses primarily based 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 strategy entails figuring out a very powerful parameters for the brand new job and updating solely these parameters throughout coaching. By doing so, PEFT can considerably scale 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 stage, that you must initialize the DPOTrainer with the next elements: the mannequin you need to practice, 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 greater LR than QLoRA paper
    max_grad_norm=0.3,                      # max gradient norm primarily based on QLoRA paper
    warmup_ratio=0.1,                       # warmup ratio primarily based 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. Greater beta means much less divergence
    "loss_type": "sigmoid"                  # The loss sort 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.practice()

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 use the saved mannequin for native inference and validation or deploy it as a SageMaker endpoint after you’ve got gained enough 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 gather 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’ve got gained enough 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 be taught extra about internet hosting your LLM utilizing DJL Serving, seek advice from Deploy massive fashions on Amazon SageMaker utilizing DJLServing and DeepSpeed mannequin parallel inference.

To deploy an LLM straight 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 referred to as serving.properties. This definition file dictates how the DJL Serving container is configured for inference.

engine = DeepSpeed
choice.tensor_parallel_degree = 1
choice.s3url = s3:///llama3-dpo-ft/modelweights
choice.hf_access_token=hf_xx1234

3. Create a customized inference file referred to 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]
    outcome = {"outputs": outputs['content']}
    return Output().add(outcome)

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, should you would like 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, keep in mind to cease your JupyterLab workspace to stop incurring extra expenses. You are able to do this by selecting Cease subsequent to your JupyterLab house. Moreover, you’ve got the choice to arrange lifecycle configuration scripts that can robotically shut down sources once they’re not in use.

In case you 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 presents 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 methods. With SageMaker Floor Reality, you may arrange workflows, handle groups, and gather constant, high-quality human suggestions.

On this submit, we confirmed the right way to improve the efficiency of Meta Llama 3 8B Instruct by fine-tuning it utilizing DPO on knowledge collected with SageMaker Floor Reality. To get began, launch SageMaker Studio and run the pocket book obtainable 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 companies.

Pranav Murthy is an AI/ML Specialist Options Architect at AWS. He focuses on serving to prospects construct, practice, deploy and migrate machine studying (ML) workloads to SageMaker. He beforehand labored within the semiconductor trade growing massive pc imaginative and prescient (CV) and pure language processing (NLP) fashions to enhance semiconductor processes. In his free time, he enjoys taking part in 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 companies. In his free time, Sundar loves touring, sports activities and having fun with out of doors actions along with his household.