Easy information to coaching Llama 2 with AWS Trainium on Amazon SageMaker

Giant language fashions (LLMs) are making a big influence within the realm of synthetic intelligence (AI). Their spectacular generative talents have led to widespread adoption throughout numerous sectors and use instances, together with content material era, sentiment evaluation, chatbot growth, and digital assistant know-how. Llama2 by Meta is an instance of an LLM supplied by AWS. Llama 2 is an auto-regressive language mannequin that makes use of an optimized transformer structure and is meant for business and analysis use in English. It is available in a variety of parameter sizes—7 billion, 13 billion, and 70 billion—in addition to pre-trained and fine-tuned variations. To study extra about Llama 2 on AWS, discuss with Llama 2 basis fashions from Meta at the moment are obtainable in Amazon SageMaker JumpStart.

Many practitioners fine-tune or pre-train these Llama 2 fashions with their very own textual content knowledge to enhance accuracy for his or her particular use case. Nevertheless, in some instances, a problem arises for practitioners: the excessive price of fine-tuning and coaching. As organizations try to push the boundaries of what LLMs can obtain, the demand for cost-effective coaching options has by no means been extra urgent. On this submit, we discover how you need to use the Neuron distributed coaching library to fine-tune, constantly pre-train, and cut back the price of coaching LLMs equivalent to Llama 2 with AWS Trainium situations on Amazon SageMaker.

AWS Trainium situations for coaching workloads

SageMaker ml.trn1 and ml.trn1n situations, powered by Trainium accelerators, are purpose-built for high-performance deep studying coaching and supply as much as 50% cost-to-train financial savings over comparable coaching optimized Amazon Elastic Compute Cloud (Amazon EC2) situations. This submit implements an answer with the ml.trn1.32xlarge Trainium occasion sort, sometimes used for coaching large-scale fashions. Nevertheless, there are additionally comparable ml.trn1n situations that provide twice as a lot networking throughput (1,600 Gbps) through Amazon Elastic Material Adapter (EFAv2). SageMaker Coaching helps the provision of ml.trn1 and ml.trn1n situations within the US East (N. Virginia) and US West (Oregon) AWS Areas, and most lately introduced common availability within the US East (Ohio) Area. These situations can be found within the listed Areas with On-Demand, Reserved, and Spot Situations, or moreover as a part of a Financial savings Plan.

For extra data on Trainium Accelerator chips, discuss with Obtain excessive efficiency with lowest price for generative AI inference utilizing AWS Inferentia2 and AWS Trainium on Amazon SageMaker. Moreover, take a look at AWS Trainium Prospects to study extra about buyer testimonials, or see Amazon EC2 Trn1 Situations for Excessive-Efficiency Mannequin Coaching are Now Accessible to dive into the accelerator highlights and specs.

Utilizing the Neuron Distributed library with SageMaker

SageMaker is a completely managed service that gives builders, knowledge scientists, and practitioners the flexibility to construct, prepare, and deploy machine studying (ML) fashions at scale. SageMaker Coaching consists of options that enhance and simplify the ML coaching expertise, together with managed infrastructure and pictures for deep studying, computerized mannequin tuning with hyperparameter optimization, and a pay-for-what-you-use billing construction. This part highlights some great benefits of utilizing SageMaker for distributed coaching with the Neuron Distributed library—particularly, the managed infrastructure, time-to-train, and cost-to-train advantages of its related resiliency and restoration options, and is a part of the AWS Neuron SDK used to run deep studying workloads on AWS Inferentia and AWS Trainum based mostly situations.

In excessive efficiency computing (HPC) clusters, equivalent to these used for deep studying mannequin coaching, {hardware} resiliency points could be a potential impediment. Though {hardware} failures whereas coaching on a single occasion could also be uncommon, points leading to stalled coaching grow to be extra prevalent as a cluster grows to tens or a whole bunch of situations. Common checkpointing helps mitigate wasted compute time, however engineering groups managing their very own infrastructure should nonetheless intently monitor their workloads and be ready to remediate a failure in any respect hours to reduce coaching downtime. The managed infrastructure of SageMaker Coaching consists of a number of resiliency options that make this monitoring and restoration course of streamlined:

• Cluster well being checks – Earlier than a coaching job begins, SageMaker runs well being checks and verifies communication on the provisioned situations. It then replaces any defective situations, if crucial, to verify the coaching script begins operating on a wholesome cluster of situations. Well being checks are at present enabled for the TRN1 occasion household in addition to P* and G* GPU-based occasion varieties.
• Computerized checkpointing – Checkpoints from a neighborhood path (/decide/ml/checkpoints by default) are mechanically copied to an Amazon Easy Storage Service (Amazon S3) location specified by the person. When coaching is restarted, SageMaker mechanically copies the beforehand saved checkpoints from the S3 location again to the native checkpoint listing to verify the coaching script can load and resume the final saved checkpoint.
• Monitoring and monitoring coaching – Within the case of a node failure, it’s vital to have the visibility of the place the failure happens. Utilizing PyTorch Neuron provides knowledge scientists the flexibility to track training progress in a TensorBoard. This lets you seize the lack of the coaching job to find out when the coaching job must be stopped to determine the convergence of the mannequin for optimum coaching.
• Constructed-in retries and cluster restore – You possibly can configure SageMaker to mechanically retry coaching jobs that fail with a SageMaker inside server error (ISE). As a part of retrying a job, SageMaker replaces any situations that encountered unrecoverable errors with contemporary situations, reboots all wholesome situations, and begins the job once more. This ends in quicker restarts and workload completion. Cluster replace is at present enabled for the TRN1 occasion household in addition to P and G GPU-based occasion varieties. Practitioners can add in their very own applicative retry mechanism across the shopper code that submits the job, to deal with different kinds of launch errors, equivalent to like exceeding your account quota.

For purchasers working with massive clusters of a whole bunch of situations for a coaching job, the resiliency and restoration options of SageMaker Coaching can cut back whole time for a mannequin to converge by as much as 20% through fewer failures and quicker restoration. This additionally allows engineering groups to watch and react to failures in any respect hours. Though SageMaker coaching jobs are appropriate for general-purpose coaching use instances with customizable configurations and integration with the broader AWS ecosystem, Amazon SageMaker HyperPod is particularly optimized for environment friendly and resilient coaching of basis fashions at scale. For extra data on SageMaker HyperPod use instances, discuss with the SageMaker HyperPod developer information.

On this submit, we use the Neuron Distributed library to constantly pre-train a Llama 2 mannequin utilizing tensor and pipeline parallelism utilizing SageMaker coaching jobs. To study extra in regards to the resiliency and restoration options of SageMaker Coaching, discuss with Coaching massive language fashions on Amazon SageMaker: Greatest practices.

Answer overview

On this answer, we use an ml.t3.medium occasion sort on a SageMaker Jupyter pocket book to course of the offered cells. We will likely be constantly pre-training our llama2-70b mannequin utilizing the trn1.32xlarge Trainium occasion. First, let’s familiarize ourselves with the methods we use to deal with the distribution of the coaching job created in our answer to contiuously pre-train our llama2-70b mannequin utilizing the Neuron distributed coaching library.

The methods used to transform the pre-trained weights within the convert_pretrained_weights.ipynb pocket book right into a .pt (PyTorch) weights file are known as pipeline parallelism and tensor parallelism:

• Pipeline parallelism includes a coaching technique that mixes parts of pipeline parallelism to optimize the coaching course of by splitting a batch or deep neural community into a number of microbatches or layers, permitting every stage employee to course of one microbatch.
• Tensor parallelism splits tensors of a neural community into a number of units. This system permits fashions with massive tensors that may’t match into the reminiscence of a single gadget.

After we convert our pre-trained weights with the previous methods in our first notebook, we comply with two separate notebooks in the identical sagemaker-trainium-examples folder. The second pocket book is Training_llama2_70b.ipynb, which walks by the continual pre-training course of by saving our checkpoint of transformed mannequin weights within the first pocket book and prepping it for inference. When this step is full, we are able to run the Convert_Nxd_to_hf.ipynb pocket book, which takes our pre-trained weights utilizing the NeuronX library and converts it right into a readable format in Hugging Face to serve inference.

Stipulations

It’s worthwhile to full some stipulations earlier than you may run the primary pocket book.

First, be sure you have created a Hugging Face access token so you may obtain the Hugging Face tokenizer for use later. After you’ve got the entry token, you must make a couple of quota improve requests for SageMaker. It’s worthwhile to request a minimal of 8 Trn1 situations ranging to a most of 32 Trn1 situations (relying on time-to-train and cost-to-train trade-offs to your use case).

On the Service Quotas console, request the next SageMaker quotas:

• Trainium situations (ml.trn1.32xlarge) for coaching job utilization: 8–32
• ml.trn1.32xlarge for coaching heat pool utilization: 8–32
• Most variety of situations per coaching job: 8–32

It could take as much as 24 hours for the quota improve to get authorized. Nevertheless, after submitting the quota improve, you may go to the sagemaker-trainium-examples GitHub repo and find the convert_pretrained_weights.ipynb file. That is the file that you just use to start the continuous pre-training course of.

Now that you just’re prepared to start the method to constantly pre-train the llama2-70b mannequin, you may convert the pre-trained weights within the subsequent part to prep the mannequin and create the checkpoint.

Getting began

Full the next steps:

1. Set up all of the required packages and libraries: SageMaker, Boto3, transformers, and datasets.

These packages just remember to can arrange your surroundings to entry your pre-trained Llama 2 mannequin, obtain your tokenizer, and get your pre-training dataset.

!pip set up -U sagemaker boto3 --quiet
!pip set up transformers datasets[s3] --quiet

2. After the packages are put in, retrieve your Hugging Face entry token, and obtain and outline your tokenizer.

The tokenizer meta-llama/Llama-2-70b-hf is a specialised tokenizer that breaks down textual content into smaller models for pure language processing. This tokenized knowledge will later be uploaded into Amazon S3 to permit for operating your coaching job.

from huggingface_hub.hf_api
import HfFolder
# Replace the entry token to obtain the tokenizer
access_token = "hf_insert-key-here"
HfFolder.save_token(access_token)

from transformers import AutoTokenizer
tokenizer_name = "meta-llama/Llama-2-70b-hf"
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
block_size = 4096

3. After following the above cells, you’ll now obtain the wikicorpus dataset from the Hugging Face dataset.
4. Tokenize the dataset with the llama-2 tokenizer that you just simply initialized.

By tokenizing the info, you’re getting ready to pre-train your Llama 2 mannequin to boost the mannequin’s efficiency to show it to the trilingual (Catalan, English, Spanish) textual content knowledge within the wikicorpus dataset to study intricate patterns and relationships within the dataset.

After the info is tokenized, run the next cell to retailer the coaching dataset to s3:

# save coaching dataset to s3
training_input_path = f's3://{sess.default_bucket()}/neuronx_distributed/knowledge'
print(f"importing coaching dataset to: {training_input_path}")
train_dataset.save_to_disk(training_input_path)

print(f"uploaded knowledge to: {training_input_path}")

The cell above makes certain that you just outline the training_input_path and have uploaded the info to your S3 bucket. You’re now prepared to start the coaching job course of.

Run the coaching job

For the coaching job, we use the trn1.32xlarge situations with every of the situations having 32 neuron cores. We use tensor parallelism and pipeline parallelism, which lets you shard the mannequin throughout Neuron cores for coaching.

The next code is the configuration for pretraining llama2-70b with trn1:

#Variety of processes per node
PROCESSES_PER_NODE = 32
# Variety of situations inside the cluster, change this if you wish to tweak the instance_count parameter
WORLD_SIZE = 32
# World batch measurement
GBS = 512
# Enter sequence size
SEQ_LEN = 4096
# Pipeline parallel diploma
PP_DEGREE = 8<br /># Tensor parallel diploma
TP_DEGREE = 8
# Information paralell measurement
DP = ((PROCESSES_PER_NODE * WORLD_SIZE / TP_DEGREE / PP_DEGREE))
# Batch measurement per mannequin duplicate
BS = ((GBS / DP))
# Quantity microbatches for pipeline execution. Setting identical as BS so every microbatch comprises a single datasample
NUM_MICROBATCHES = BS
# Variety of whole steps for which to coach mannequin. This quantity must be adjusted to the step quantity when the loss operate is approaching convergence.
MAX_STEPS = 1500
# Timeout in seconds for coaching. After this period of time Amazon SageMaker terminates the job no matter its present standing.
MAX_RUN = 2 * (24 * 60 * 60)

Now you may outline the hyperparameters for coaching. Observe that adjusting these parameters based mostly on {hardware} capabilities, dataset traits, and convergence necessities can considerably influence coaching efficiency and effectivity.

The next is the code for the hyperparameters:

hyperparameters = {}
hyperparameters["train_batch_size"] = int(BS)
hyperparameters["use_meta_device_init"] = 1
hyperparameters["training_dir"] = "/decide/ml/enter/knowledge/prepare" # path the place sagemaker uploads the coaching knowledge
hyperparameters["training_config"] = "config.json" # config file containing llama 70b configuration , change this for tweaking the variety of parameters.

hyperparameters["max_steps"] = MAX_STEPS
hyperparameters["seq_len"] = SEQ_LEN
hyperparameters["pipeline_parallel_size"] = PP_DEGREE
hyperparameters["tensor_parallel_size"] = TP_DEGREE
hyperparameters["num_microbatches"] = int(NUM_MICROBATCHES)
hyperparameters["lr"] = 0.00015
hyperparameters["min_lr"] = 1e-05
hyperparameters["beta1"] = 0.9
hyperparameters["beta2"] = 0.95
hyperparameters["weight_decay"] = 0.1
hyperparameters["warmup_steps"] = 2000
hyperparameters["constant_steps"] = 0
hyperparameters["use_zero1_optimizer"] = 1
hyperparameters["tb_dir"] = "/decide/ml/checkpoints/tensorboard" # The tensorboard logs will likely be saved right here and finally pushed to S3.

Now you specify the Docker picture that will likely be used to coach the mannequin on Trainium:

docker_image = f"763104351884.dkr.ecr.{region_name}.amazonaws.com/pytorch-training-neuronx:1.13.1-neuronx-py310-sdk2.18.0-ubuntu20.04"

The picture we outlined is designed for PyTorch coaching with Neuron optimizations. This picture is configured to work with PyTorch, utilizing Neuron SDK model 2.18.0 for enhanced efficiency and effectivity on Trn1 situations outfitted with AWS Trainium chips. This picture can also be appropriate with Python 3.10, indicated by the py310, and is predicated on Ubuntu 20.04.

Previous to beginning your coaching job, you must configure it by defining all crucial variables. You accomplish that by defining the coaching job title, checkpoint listing, and cache listing:

import time
# Outline Coaching Job Title
job_name = f'llama-neuron-{time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())}'
# Outline checkpoint listing that comprises the weights and different related knowledge for the skilled mannequin
checkpoint_s3_uri = "s3://" + sagemaker_session_bucket + "/neuron_llama_experiment"
checkpoint_dir="/decide/ml/checkpoints"</p><p>
In [ ]:
# Outline neuron chache listing
cache_dir = "/decide/ml/checkpoints/neuron_cache"

The parameters allow you to do the next:

• The coaching job lets you determine and observe particular person coaching jobs based mostly on timestamps
• The checkpoint listing specifies the S3 URI the place the checkpoint knowledge, weights, and different data are saved for the skilled mannequin
• The cache listing helps optimize the coaching course of by storing and reusing beforehand calculated values, from the checkpoint listing, decreasing redundancy and bettering effectivity
• The surroundings variables guarantee that the coaching job is optimally configured and settings are tailor-made to allow environment friendly and efficient coaching utilizing options like RDMA, optimized reminiscence allocation, fused operations, and Neuron-specific gadget optimizations

After you’ve got outlined your coaching job and configured all directories and surroundings variables for an optimum coaching pipeline, you now arrange your PyTorch estimator to start the coaching job on SageMaker. A SageMaker estimator is a high-level interface that handles the end-to-end SageMaker coaching and deployment duties.

The entry_point is specified because the Python script run_llama_nxd.py. We use the instance_type ml.trn1.32xlarge, the occasion rely is 32 (which was beforehand outlined as a world variable within the configuration code), and input_mode is about to FastFile. Quick File mode in SageMaker streams knowledge from Amazon S3 on demand, which optimizes knowledge loading efficiency by fetching knowledge as wanted, decreasing total useful resource consumption. For extra data on enter, discuss with Entry Coaching Information.

from sagemaker.pytorch import PyTorch

# Deal with end-to-end Amazon SageMaker coaching and deployment duties.
pt_estimator = PyTorch(<br />entry_point="run_llama_nxd.py",
source_dir="./scripts",<br />instance_type="ml.trn1.32xlarge",
image_uri=docker_image,<br />instance_count=WORLD_SIZE,
max_run=MAX_RUN,
hyperparameters=hyperparameters,
function=function,
base_job_name=job_name,
surroundings=env,
input_mode="FastFile",
disable_output_compression=True,
keep_alive_period_in_seconds=600, # that is added to allow heat pool functionality
checkpoint_s3_uri=checkpoint_s3_uri,
checkpoint_local_path=checkpoint_dir,
distribution={"torch_distributed": {"enabled": True}} # allow torchrun
)

Lastly, you can begin the coaching job with the SageMaker match() methodology, which trains the mannequin based mostly on the outlined hyperparameters:

# Begin coaching job
pt_estimator.match({"prepare": training_input_path})

You have got efficiently began the method to constantly pre-train a llama2-70b mannequin by changing pre-trained weights with tokenized knowledge utilizing SageMaker coaching on Trainium situations.

Steady pre-training

After following the stipulations, finishing the offered pocket book, and changing the pre-trained weights as a checkpoint, now you can start the continuous pre-training course of, utilizing the checkpoint as a degree of reference to pre-train the llama2-70b mannequin. The methods used to transform the pre-trained weights within the convert_pretrained_weights.ipynb pocket book right into a .pt (PyTorch) weights file are known as pipeline parallelism and tensor parallelism.

To start the continual pre-training course of, comply with the Training_llama2_70b.ipynb file within the sagemaker-trainium-examples repo.

Given the big measurement of the llama2-70b mannequin, you must convert the pre-trained weights right into a extra environment friendly and useable format (.pt). You are able to do so by defining the hyperparameters in your configuration to retailer transformed weights and checkpoints. The next are the hyperparameters:

# Use the sagemaker s3 checkpoints mechanism since we want learn/write entry to the paths.
hyperparameters["output_dir"] = "/decide/ml/checkpoints/llama70b_weights"
hyperparameters["checkpoint-dir"] = '/decide/ml/checkpoints'<br />hyperparameters["n_layers"] = 80
hyperparameters["convert_from_full_model"] = ""

If you happen to take a look at the hyperparameters, the output_dir is used as a reference for pre-training. In case you are at this cell, it’s best to have already adopted the Training_llama2_70b.ipynb pocket book and gone by the method of establishing your SageMaker shopper and Docker picture, and getting ready the pre-trained weights for pre-training. You’re now able to carry out the continual pre-training course of on the llama2-70b mannequin.

We use the next parameters to take the pre-trained weights saved in output_dir within the convert_pretrained_weights.ipynb file to be reused constantly for pre-training:

hyperparameters["checkpoint_dir"] = "/decide/ml/checkpoints/checkpts"
hyperparameters["checkpoint_freq"] = 10
hyperparameters["num_kept_checkpoint"] = 1
hyperparameters["use_zero1_optimizer"] = 1
hyperparameters["save_load_xser"] = 0
hyperparameters["pretrained_weight_dir"] = "/decide/ml/checkpoints/llama70b_weights"

After these hyperparameters are carried out, you may run the remainder of the pocket book cells to finish the continual pre-training course of. After the SageMaker estimator has accomplished the coaching job, you may find the brand new checkpoint within the S3 checkpoint listing containing the weights. Now you can find the convert_Nxd_to_hf.ipynb file to get the checkpoint prepared for inferencing.

Convert the Neuron Distributed checkpoint for inferencing

Checkpoints play an important function within the context of distributed coaching with the NeuronX library as a result of it has checkpoint compatibility with Hugging Face Transformers. You will get the coaching job output prepared for inferencing by taking the coaching job that’s saved as a NeuronX distributed checkpoint and changing the weights into .pt weights information.

To transform the checkpoints to Hugging Face format utilizing NeuronX, you first want to save lots of the S3 nxd_checkpoint_path listing:

# S3 checkpoint listing that comprises the weights and different related knowledge from the continual pre-trained mannequin
checkpoint_s3_uri = "<pre-training-checkpoint-s3-uri>"
nxd_checkpoint_path = f"s3://{checkpoint_s3_uri}/neuronx_llama_experiment/checkpts/step10/mannequin/"
# Checkpoint is saved as a part of Pocket book 2

After you save the checkpoint within the nxd_checkpoint_path listing, it can save you your hyperparameters and configure your SageMaker estimator, which makes certain the pre-training course of can start. Now you can run the match() operate inside the estimator to transform the pre-trained weights right into a checkpoint for inferencing with the next cell:

# Begin SageMaker job
estimator.match({"checkpoint": nxd_checkpoint_path})

Abstract

You have got efficiently carried out steady pre-training on a llama2-70b mannequin by changing your pre-trained weights and checkpoint for use to serve inference utilizing the Neuron SDK and Trainium situations. By following the answer on this submit, it’s best to now know how you can configure a pipeline for steady pre-training of an LLM utilizing SageMaker and Trainium accelerator chips.

For extra data on how you can use Trainium to your workloads, discuss with the Neuron SDK documentation or attain out on to the workforce. We worth buyer suggestions and are all the time trying to have interaction with ML practitioners and builders. Be at liberty to depart feedback or questions within the feedback part.

Concerning the authors

Marco Punio is a Options Architect targeted on generative AI technique, utilized AI options and conducting analysis to assist prospects hyperscale on AWS. He’s a certified technologist with a ardour for machine studying, synthetic intelligence, and mergers & acquisitions. Marco is predicated in Seattle, WA and enjoys writing, studying, exercising, and constructing purposes in his free time.

Armando Diaz is a Options Architect at AWS. He focuses on generative AI, AI/ML, and Information Analytics. At AWS, Armando helps prospects integrating cutting-edge generative AI capabilities into their techniques, fostering innovation and aggressive benefit. When he’s not at work, he enjoys spending time along with his spouse and household, climbing, and touring the world.

Arun Kumar Lokanatha is a Senior ML Options Architect with the Amazon SageMaker Service workforce. He focuses on serving to prospects construct, prepare, and migrate ML manufacturing workloads to SageMaker at scale. He focuses on deep studying, particularly within the space of NLP and CV. Exterior of labor, he enjoys operating and climbing.

Robert Van Dusen is a Senior Product Supervisor with Amazon SageMaker. He leads frameworks, compilers, and optimization methods for deep studying coaching.

Niithiyn Vijeaswaran is a Options Architect at AWS. His space of focus is generative AI and AWS AI Accelerators. He holds a Bachelor’s diploma in Laptop Science and Bioinformatics. Niithiyn works intently with the Generative AI GTM workforce to allow AWS prospects on a number of fronts and speed up their adoption of generative AI. He’s an avid fan of the Dallas Mavericks and enjoys gathering sneakers.

Rohit Talluri is a Generative AI GTM Specialist (Tech BD) at Amazon Internet Providers (AWS). He’s partnering with prime generative AI mannequin builders, strategic prospects, key AI/ML companions, and AWS Service Groups to allow the subsequent era of synthetic intelligence, machine studying, and accelerated computing on AWS. He was beforehand an Enterprise Options Architect, and the World Options Lead for AWS Mergers & Acquisitions Advisory.

Sebastian Bustillo is a Options Architect at AWS. He focuses on AI/ML applied sciences with a profound ardour for generative AI and compute accelerators. At AWS, he helps prospects unlock enterprise worth by generative AI. When he’s not at work, he enjoys brewing an ideal cup of specialty espresso and exploring the world along with his spouse.