Advantageous-tune Meta Llama 3.1 fashions utilizing torchtune on Amazon SageMaker

This submit is co-written with Meta’s PyTorch group.

In as we speak’s quickly evolving AI panorama, companies are continually searching for methods to make use of superior giant language fashions (LLMs) for his or her particular wants. Though basis fashions (FMs) supply spectacular out-of-the-box capabilities, true aggressive benefit usually lies in deep mannequin customization by way of fine-tuning. Nonetheless, fine-tuning LLMs for complicated duties usually requires superior AI experience to align and optimize them successfully. Recognizing this problem, Meta developed torchtune, a PyTorch-native library that simplifies authoring, fine-tuning, and experimenting with LLMs, making it extra accessible to a broader vary of customers and functions.

On this submit, AWS collaborates with Meta’s PyTorch group to showcase how you should utilize Meta’s torchtune library to fine-tune Meta Llama-like architectures whereas utilizing a fully-managed surroundings offered by Amazon SageMaker Coaching. We show this by way of a step-by-step implementation of mannequin fine-tuning, inference, quantization, and analysis. We carry out the steps on a Meta Llama 3.1 8B mannequin using the LoRA fine-tuning technique on a single p4d.24xlarge employee node (offering 8 Nvidia A100 GPUs).

Earlier than we dive into the step-by-step information, we first explored the efficiency of our technical stack by fine-tuning a Meta Llama 3.1 8B mannequin throughout varied configurations and occasion sorts.

As could be seen within the following chart, we discovered {that a} single p4d.24xlarge delivers 70% greater efficiency than two g5.48xlarge situations (every with 8 NVIDIA A10 GPUs) at nearly 47% lowered value. We due to this fact have optimized the instance on this submit for a p4d.24xlarge configuration. Nonetheless, you can use the identical code to run single-node or multi-node coaching on completely different occasion configurations by altering the parameters handed to the SageMaker estimator. You might additional optimize the time for coaching within the following graph by utilizing a SageMaker managed heat pool and accessing pre-downloaded fashions utilizing Amazon Elastic File System (Amazon EFS).

Challenges with fine-tuning LLMs

Generative AI fashions supply many promising enterprise use circumstances. Nonetheless, to take care of factual accuracy and relevance of those LLMs to particular enterprise domains, fine-tuning is required. As a result of rising variety of mannequin parameters and the growing context size of recent LLMs, this course of is reminiscence intensive. To deal with these challenges, fine-tuning methods like LoRA (Low-Rank Adaptation) and QLoRA (Quantized Low-Rank Adaptation) restrict the variety of trainable parameters by including low-rank parallel constructions to the transformer layers. This lets you prepare LLMs even on methods with low reminiscence availability like commodity GPUs. Nonetheless, this results in an elevated complexity as a result of new dependencies must be dealt with and coaching recipes and hyperparameters must be tailored to the brand new strategies.

What companies want as we speak is user-friendly coaching recipes for these common fine-tuning strategies, which give abstractions to the end-to-end tuning course of, addressing the frequent pitfalls in probably the most opinionated means.

How does torchtune helps?

torchtune is a PyTorch-native library that goals to democratize and streamline the fine-tuning course of for LLMs. By doing so, it makes it simple for researchers, builders, and organizations to adapt these highly effective LLMs to their particular wants and constraints. It supplies coaching recipes for quite a lot of fine-tuning strategies, which could be configured by way of YAML information. The recipes implement frequent fine-tuning strategies (full-weight, LoRA, QLoRA) in addition to different frequent duties like inference and analysis. They robotically apply a set of essential options (FSDP, activation checkpointing, gradient accumulation, blended precision) and are particular to a given mannequin household (corresponding to Meta Llama 3/3.1 or Mistral) in addition to compute surroundings (single-node vs. multi-node).

Moreover, torchtune integrates with main libraries and frameworks like Hugging Face datasets, EleutherAI’s Eval Harness, and Weights & Biases. This helps handle the necessities of the generative AI fine-tuning lifecycle, from knowledge ingestion and multi-node fine-tuning to inference and analysis. The next diagram exhibits a visualization of the steps we describe on this submit.

Confer with the set up directions and PyTorch documentation to be taught extra about torchtune and its ideas.

Answer overview

This submit demonstrates using SageMaker Coaching for operating torchtune recipes by way of task-specific coaching jobs on separate compute clusters. SageMaker Coaching is a complete, absolutely managed ML service that permits scalable mannequin coaching. It supplies versatile compute useful resource choice, help for customized libraries, a pay-as-you-go pricing mannequin, and self-healing capabilities. By managing workload orchestration, well being checks, and infrastructure, SageMaker helps cut back coaching time and whole price of possession.

The answer structure incorporates the next key parts to reinforce safety and effectivity in fine-tuning workflows:

• Safety enhancement – Coaching jobs are run inside personal subnets of your digital personal cloud (VPC), considerably enhancing the safety posture of machine studying (ML) workflows.
• Environment friendly storage resolution – Amazon EFS is used to speed up mannequin storage and entry throughout varied phases of the ML workflow.
• Customizable surroundings – We use customized containers in coaching jobs. The help in SageMaker for customized containers permits you to package deal all mandatory dependencies, specialised frameworks, and libraries right into a single artifact, offering full management over your ML surroundings.

The next diagram illustrates the answer structure. Customers provoke the method by calling the SageMaker management airplane by way of APIs or command line interface (CLI) or utilizing the SageMaker SDK for every particular person step. In response, SageMaker spins up coaching jobs with the requested quantity and sort of compute situations to run particular duties. Every step outlined within the diagram accesses torchtune recipes from an Amazon Easy Storage Service (Amazon S3) bucket and makes use of Amazon EFS to avoid wasting and entry mannequin artifacts throughout completely different levels of the workflow.

By decoupling each torchtune step, we obtain a stability between flexibility and integration, permitting for each impartial execution of steps and the potential for automating this course of utilizing seamless pipeline integration.

On this use case, we fine-tune a Meta Llama 3.1 8B mannequin with LoRA. Subsequently, we run mannequin inference, and optionally quantize and consider the mannequin utilizing torchtune and SageMaker Coaching.

Recipes, configs, datasets, and immediate templates are fully configurable and can help you align torchtune to your necessities. To show this, we use a customized immediate template on this use case and mix it with the open supply dataset Samsung/samsum from the Hugging Face hub.

We fine-tune the mannequin utilizing torchtune’s multi system LoRA recipe (lora_finetune_distributed) and use the SageMaker personalized model of Meta Llama 3.1 8B default config (llama3_1/8B_lora).

Conditions

You could full the next stipulations earlier than you’ll be able to run the SageMaker Jupyter notebooks:

1. Create a Hugging Face entry token to get entry to the gated repo meta-llama/Meta-Llama-3.1-8B on Hugging Face.
2. Create a Weights & Biases API key to entry the Weights & Biases dashboard for logging and monitoring
3. Request a SageMaker service quota for 1x ml.p4d.24xlarge and 1xml.g5.2xlarge.
4. Create an AWS Id and Entry Administration (IAM) position with managed insurance policies AmazonSageMakerFullAccess, AmazonEC2FullAccess, AmazonElasticFileSystemFullAccess, and AWSCloudFormationFullAccess to provide required entry to SageMaker to run the examples. (That is for demonstration functions. You must modify this to your particular safety necessities for manufacturing.)
5. Create an Amazon SageMaker Studio area (see Fast setup to Amazon SageMaker) to entry Jupyter notebooks with the previous position. Confer with the directions to set permissions for Docker construct.
6. Log in to the pocket book console and clone the GitHub repo:

$ git clone https://github.com/aws-samples/sagemaker-distributed-training-workshop.git
$ cd sagemaker-distributed-training-workshop/13-torchtune

7. Run the pocket book ipynb to arrange VPC and Amazon EFS utilizing an AWS CloudFormation stack.

Evaluate torchtune configs

The next determine illustrates the steps in our workflow.

You possibly can lookup the torchtune configs in your use case by instantly utilizing the tune CLI.For this submit, we offer modified config information aligned with SageMaker listing path’s construction:

sh-4.2$ cd config/
sh-4.2$ ls -ltr
-rw-rw-r-- 1 ec2-user ec2-user 1151 Aug 26 18:34 config_l3.1_8b_gen_orig.yaml
-rw-rw-r-- 1 ec2-user ec2-user 1172 Aug 26 18:34 config_l3.1_8b_gen_trained.yaml
-rw-rw-r-- 1 ec2-user ec2-user  644 Aug 26 18:49 config_l3.1_8b_quant.yaml
-rw-rw-r-- 1 ec2-user ec2-user 2223 Aug 28 14:53 config_l3.1_8b_lora.yaml
-rw-rw-r-- 1 ec2-user ec2-user 1223 Sep  4 14:28 config_l3.1_8b_eval_trained.yaml
-rw-rw-r-- 1 ec2-user ec2-user 1213 Sep  4 14:29 config_l3.1_8b_eval_original.yaml

torchtune makes use of these config information to pick out and configure the parts (assume fashions and tokenizers) throughout the execution of the recipes.

Construct the container

As a part of our instance, we create a customized container to offer customized libraries like torch nightlies and torchtune. Full the next steps:

sh-4.2$ cat Dockerfile
# Set the default worth for the REGION construct argument
ARG REGION=us-west-2
# SageMaker PyTorch picture for TRAINING
FROM ${ACCOUNTID}.dkr.ecr.${REGION}.amazonaws.com/pytorch-training:2.3.0-gpu-py311-cu121-ubuntu20.04-sagemaker
# Uninstall current PyTorch packages
RUN pip uninstall torch torchvision transformer-engine -y
# Set up newest launch of PyTorch and torchvision
RUN pip set up --force-reinstall torch==2.4.1 torchao==0.4.0 torchvision==0.19.1

Run the 1_build_container.ipynb pocket book till the next command to push this file to your ECR repository:

!sm-docker construct . --repository speed up:newest

sm-docker is a CLI instrument designed for constructing Docker photos in SageMaker Studio utilizing AWS CodeBuild. We set up the library as a part of the pocket book.

Subsequent, we are going to run the 2_torchtune-llama3_1.ipynb pocket book for all fine-tuning workflow duties.

For each activity, we assessment three artifacts:

• torchtune configuration file
• SageMaker activity config with compute and torchtune recipe particulars
• SageMaker activity output

Run the fine-tuning activity

On this part, we stroll by way of the steps to run and monitor the fine-tuning activity.

Run the fine-tuning job

The next code exhibits a shortened torchtune recipe configuration highlighting just a few key parts of the file for a fine-tuning job:

• Mannequin part together with LoRA rank configuration
• Meta Llama 3 tokenizer to tokenize the information
• Checkpointer to learn and write checkpoints
• Dataset part to load the dataset

sh-4.2$ cat config_l3.1_8b_lora.yaml
# Mannequin Arguments
mannequin:
  _component_: torchtune.fashions.llama3_1.lora_llama3_1_8b
  lora_attn_modules: ['q_proj', 'v_proj']
  lora_rank: 8
  lora_alpha: 16

# Tokenizer
tokenizer:
  _component_: torchtune.fashions.llama3.llama3_tokenizer
  path: /choose/ml/enter/knowledge/mannequin/hf-model/unique/tokenizer.mannequin

checkpointer:
  _component_: torchtune.utils.FullModelMetaCheckpointer
  checkpoint_files: [
    consolidated.00.pth
  ]
  …

# Dataset and Sampler
dataset:
  _component_: torchtune.datasets.samsum_dataset
  train_on_input: True
batch_size: 13

# Coaching
epochs: 1
gradient_accumulation_steps: 2

... and extra ...

We use Weights & Biases for logging and monitoring our coaching jobs, which helps us observe our mannequin’s efficiency:

metric_logger:
_component_: torchtune.utils.metric_logging.WandBLogger
…

Subsequent, we outline a SageMaker activity that will probably be handed to our utility operate within the script create_pytorch_estimator. This script creates the PyTorch estimator with all of the outlined parameters.

Within the activity, we use the lora_finetune_distributed torchrun recipe with config config-l3.1-8b-lora.yaml on an ml.p4d.24xlarge occasion. Be sure to obtain the bottom mannequin from Hugging Face earlier than it’s fine-tuned utilizing the use_downloaded_model parameter. The image_uri parameter defines the URI of the customized container.

sagemaker_tasks={
    "fine-tune":{
        "hyperparameters":{
            "tune_config_name":"config-l3.1-8b-lora.yaml",
            "tune_action":"fine-tune",
            "use_downloaded_model":"false",
            "tune_recipe":"lora_finetune_distributed"
            },
        "instance_count":1,
        "instance_type":"ml.p4d.24xlarge",        
        "image_uri":"<accountid>.dkr.ecr.<area>.amazonaws.com/speed up:newest"
    }
    ... and extra ...
}

To create and run the duty, run the next code:

Activity="fine-tune"
estimator=create_pytorch_estimator(**sagemaker_tasks[Task])
execute_task(estimator)

The next code exhibits the duty output and reported standing:

# Refer-Output

2024-08-16 17:45:32 Beginning - Beginning the coaching job...
...
...

1|140|Loss: 1.4883038997650146:  99%|█████████▉| 141/142 [06:26<00:02,  2.47s/it]
1|141|Loss: 1.4621509313583374:  99%|█████████▉| 141/142 [06:26<00:02,  2.47s/it]

Coaching accomplished with code: 0
2024-08-26 14:19:09,760 sagemaker-training-toolkit INFO     Reporting coaching SUCCESS

The ultimate mannequin is saved to Amazon EFS, which makes it accessible with out obtain time penalties.

Monitor the fine-tuning job

You possibly can monitor varied metrics corresponding to loss and studying price in your coaching run by way of the Weights & Biases dashboard. The next figures present the outcomes of the coaching run the place we tracked GPU utilization, GPU reminiscence utilization, and loss curve.

For the next graph, to optimize reminiscence utilization, torchtune makes use of solely rank 0 to initially load the mannequin into CPU reminiscence. rank 0 due to this fact will probably be answerable for loading the mannequin weights from the checkpoint.

The instance is optimized to make use of GPU reminiscence to its most capability. Rising the batch dimension additional will result in CUDA out-of-memory (OOM) errors.

The run took about 13 minutes to finish for one epoch, ensuing within the loss curve proven within the following graph.

Run the mannequin technology activity

Within the subsequent step, we use the beforehand fine-tuned mannequin weights to generate the reply to a pattern immediate and examine it to the bottom mannequin.

The next code exhibits the configuration of the generate recipe config_l3.1_8b_gen_trained.yaml. The next are key parameters:

• FullModelMetaCheckpointer – We use this to load the educated mannequin checkpoint meta_model_0.pt from Amazon EFS
• CustomTemplate.SummarizeTemplate – We use this to format the immediate for inference

# torchtune - educated mannequin technology config - config_l3.1_8b_gen_trained.yaml
mannequin:
  _component_: torchtune.fashions.llama3_1.llama3_1_8b
  
checkpointer:
  _component_: torchtune.utils.FullModelMetaCheckpointer
  checkpoint_dir: /choose/ml/enter/knowledge/mannequin/
  checkpoint_files: [
    meta_model_0.pt
  ]
  …

# Era arguments; defaults taken from gpt-fast
instruct_template: CustomTemplate.SummarizeTemplate

... and extra ...

Subsequent, we configure the SageMaker activity to run on a single ml.g5.2xlarge occasion:

immediate=r'{"dialogue":"Amanda: I baked  cookies. Would you like some?rnJerry: Certain rnAmanda: I'll convey you tomorrow :-)"}'

sagemaker_tasks={
    "generate_inference_on_trained":{
        "hyperparameters":{
            "tune_config_name":"config_l3.1_8b_gen_trained.yaml ",
            "tune_action":"generate-trained",
            "use_downloaded_model":"true",
            "immediate":json.dumps(immediate)
            },
        "instance_count":1,
        "instance_type":"ml.g5.2xlarge",
 "image_uri":"<accountid>.dkr.ecr.<area>.amazonaws.com/speed up:newest"
    }
}

Within the output of the SageMaker activity, we see the mannequin abstract output and a few stats like tokens per second:

#Refer- Output
...
Amanda: I baked  cookies. Would you like some?rnJerry: Certain rnAmanda: I'll convey you tomorrow :-)

Abstract:
Amanda baked cookies. She's going to convey some to Jerry tomorrow.

INFO:torchtune.utils.logging:Time for inference: 1.71 sec whole, 7.61 tokens/sec
INFO:torchtune.utils.logging:Reminiscence used: 18.32 GB

... and extra ...

We are able to generate inference from the unique mannequin utilizing the unique mannequin artifact consolidated.00.pth:

# torchtune - educated unique technology config - config_l3.1_8b_gen_orig.yaml
…  
checkpointer:
  _component_: torchtune.utils.FullModelMetaCheckpointer
  checkpoint_dir: /choose/ml/enter/knowledge/mannequin/hf-model/unique/
  checkpoint_files: [
    consolidated.00.pth
  ]
  
... and extra ...

The next code exhibits the comparability output from the bottom mannequin run with the SageMaker activity (generate_inference_on_original). We are able to see that the fine-tuned mannequin is performing subjectively higher than the bottom mannequin by additionally mentioning that Amanda baked the cookies.

# Refer-Output 
---
Abstract:
Jerry tells Amanda he needs some cookies. Amanda says she is going to convey him some cookies tomorrow.

... and extra ...

Run the mannequin quantization activity

To hurry up the inference and reduce the mannequin artifact dimension, we are able to apply post-training quantization. torchtune depends on torchao for post-training quantization.

We configure the recipe to make use of Int8DynActInt4WeightQuantizer, which refers to int8 dynamic per token activation quantization mixed with int4 grouped per axis weight quantization. For extra particulars, discuss with the torchao implementation.

# torchtune mannequin quantization config - config_l3.1_8b_quant.yaml
mannequin:
  _component_: torchtune.fashions.llama3_1.llama3_1_8b

checkpointer:
  _component_: torchtune.utils.FullModelMetaCheckpointer
  …

quantizer:
  _component_: torchtune.utils.quantization.Int8DynActInt4WeightQuantizer
  groupsize: 256

We once more use a single ml.g5.2xlarge occasion and use SageMaker heat pool configuration to hurry up the spin-up time for the compute nodes:

sagemaker_tasks={
"quantize_trained_model":{
        "hyperparameters":{
            "tune_config_name":"config_l3.1_8b_quant.yaml",
            "tune_action":"run-quant",
            "use_downloaded_model":"true"
            },
        "instance_count":1,
        "instance_type":"ml.g5.2xlarge",
        "image_uri":"<accountid>.dkr.ecr.<area>.amazonaws.com/speed up:newest"
    }
}

Within the output, we see the situation of the quantized mannequin and the way a lot reminiscence we saved as a result of course of:

#Refer-Output
...

linear: layers.31.mlp.w1, in=4096, out=14336
linear: layers.31.mlp.w2, in=14336, out=4096
linear: layers.31.mlp.w3, in=4096, out=14336
linear: output, in=4096, out=128256
INFO:torchtune.utils.logging:Time for quantization: 7.40 sec
INFO:torchtune.utils.logging:Reminiscence used: 22.97 GB
INFO:torchtune.utils.logging:Mannequin checkpoint of dimension 8.79 GB saved to /choose/ml/enter/knowledge/mannequin/quantized/meta_model_0-8da4w.pt

... and extra ...

You possibly can run mannequin inference on the quantized mannequin meta_model_0-8da4w.pt by updating the inference-specific configurations.

Run the mannequin analysis activity

Lastly, let’s consider our fine-tuned mannequin in an goal method by operating an analysis on the validation portion of our dataset.

torchtune integrates with EleutherAI’s analysis harness and supplies the eleuther_eval recipe.

For our analysis, we use a customized activity for the analysis harness to guage the dialogue summarizations utilizing the rouge metrics.

The recipe configuration factors the analysis harness to our customized analysis activity:

# torchtune educated mannequin analysis config - config_l3.1_8b_eval_trained.yaml

mannequin:
...

include_path: "/choose/ml/enter/knowledge/config/duties"
duties: ["samsum"]
...

The next code is the SageMaker activity that we run on a single ml.p4d.24xlarge occasion:

sagemaker_tasks={
"evaluate_trained_model":{
        "hyperparameters":{
            "tune_config_name":"config_l3.1_8b_eval_trained.yaml",
            "tune_action":"run-eval",
            "use_downloaded_model":"true",
            },
        "instance_count":1,
        "instance_type":"ml.p4d.24xlarge",
    }
}

Run the mannequin analysis on ml.p4d.24xlarge:

Activity="evaluate_trained_model"
estimator=create_pytorch_estimator(**sagemaker_tasks[Task])
execute_task(estimator)

The next tables present the duty output for the fine-tuned mannequin in addition to the bottom mannequin.

The next output is for the fine-tuned mannequin.

The next output is for the bottom mannequin.

Our fine-tuned mannequin achieves an enchancment of roughly 46% on the summarization activity, which is roughly 12 factors higher than the baseline.

Clear up

Full the next steps to wash up your assets:

1. Delete any unused SageMaker Studio assets.
2. Optionally, delete the SageMaker Studio area.
3. Delete the CloudFormation stack to delete the VPC and Amazon EFS assets.

Conclusion

On this submit, we mentioned how one can fine-tune Meta Llama-like architectures utilizing varied fine-tuning methods in your most popular compute and libraries, utilizing customized dataset immediate templates with torchtune and SageMaker. This structure provides you a versatile means of operating fine-tuning jobs which might be optimized for GPU reminiscence and efficiency. We demonstrated this by way of fine-tuning a Meta Llama3.1 mannequin utilizing P4 and G5 situations on SageMaker and used observability instruments like Weights & Biases to watch loss curve, in addition to CPU and GPU utilization.

We encourage you to make use of SageMaker coaching capabilities and Meta’s torchtune library to fine-tune Meta Llama-like architectures in your particular enterprise use circumstances. To remain knowledgeable about upcoming releases and new options, discuss with the torchtune GitHub repo and the official Amazon SageMaker coaching documentation .

Particular because of Kartikay Khandelwal (Software program Engineer at Meta), Eli Uriegas (Engineering Supervisor at Meta), Raj Devnath (Sr. Product Supervisor Technical at AWS) and Arun Kumar Lokanatha (Sr. ML Answer Architect at AWS) for his or her help to the launch of this submit.

In regards to the Authors

Kanwaljit Khurmi is a Principal Options Architect at Amazon Net Providers. He works with AWS clients to offer steerage and technical help, serving to them enhance the worth of their options when utilizing AWS. Kanwaljit focuses on serving to clients with containerized and machine studying functions.

Roy Allela is a Senior AI/ML Specialist Options Architect at AWS.He helps AWS clients—from small startups to giant enterprises—prepare and deploy giant language fashions effectively on AWS.

Matthias Reso is a Companion Engineer at PyTorch engaged on open supply, high-performance mannequin optimization, distributed coaching (FSDP), and inference. He’s a co-maintainer of llama-recipes and TorchServe.

Trevor Harvey is a Principal Specialist in Generative AI at Amazon Net Providers (AWS) and an AWS Licensed Options Architect – Skilled. He serves as a voting member of the PyTorch Basis Governing Board, the place he contributes to the strategic development of open-source deep studying frameworks. At AWS, Trevor works with clients to design and implement machine studying options and leads go-to-market methods for generative AI providers.