Environment friendly Pre-training of Llama 3-like mannequin architectures utilizing torchtitan on Amazon SageMaker

This put up is co-written with Much less Wright and Wei Feng from Meta

Pre-training massive language fashions (LLMs) is step one in creating highly effective AI programs that may perceive and generate human-like textual content. By exposing fashions to huge quantities of various knowledge, pre-training lays the groundwork for LLMs to study basic language patterns, world data, and reasoning capabilities. This foundational course of permits LLMs to carry out a variety of duties with out task-specific coaching, making them extremely versatile and adaptable. Pre-training is important for constructing a powerful base of data, which might then be refined and specialised by means of fine-tuning, switch studying, or few-shot studying approaches.

On this put up, we collaborate with the crew engaged on PyTorch at Meta to showcase how the torchtitan library accelerates and simplifies the pre-training of Meta Llama 3-like mannequin architectures. We showcase the important thing options and capabilities of torchtitan reminiscent of FSDP2, torch.compile integration, and FP8 assist that optimize the coaching effectivity. We pre-train a Meta Llama 3 8B mannequin structure utilizing torchtitan on Amazon SageMaker on p5.48xlarge cases, every geared up with 8 Nvidia H100 GPUs. We display a 38.23% efficiency speedup within the coaching throughput in comparison with the baseline with out making use of the optimizations (as proven within the following determine). Amazon SageMaker Mannequin Coaching reduces the time and price to coach and tune machine studying (ML) fashions at scale with out the necessity to handle infrastructure. You possibly can reap the benefits of the highest-performing ML compute infrastructure at present obtainable, and SageMaker can routinely scale infrastructure up or down, from one to hundreds of GPUs.

To study extra, you will discover our full code pattern on GitHub.

Introduction to torchtitan

torchtitan is a reference structure for large-scale LLM coaching utilizing native PyTorch. It goals to showcase PyTorch’s newest distributed coaching options in a clear, minimal code base. The library is designed to be easy to know, use, and lengthen for various coaching functions, with minimal modifications required to the mannequin code when making use of varied parallel processing methods.

torchtitan presents a number of key options, together with FSDP2 with per-parameter sharding, tensor parallel processing, selective layer and operator activation checkpointing, and distributed checkpointing. It helps pre-training of Meta Llama 3-like and Llama 2-like mannequin architectures of varied sizes and contains configurations for a number of datasets. The library supplies easy configuration by means of TOML recordsdata and presents efficiency monitoring by means of TensorBoard. Within the following sections, we spotlight among the key options of torchtitan.

Transitioning from FSDP1 to FSDP2

FSDP1 and FSDP2 are two approaches to completely sharded knowledge parallel coaching. FSDP1 makes use of flat-parameter sharding, which flattens all parameters to 1D, concatenates them right into a single tensor, pads it, after which chunks it throughout staff. This methodology presents bounded padding and environment friendly unsharded storage, however may not at all times enable optimum sharding for particular person parameters. FSDP2, then again, represents sharded parameters as DTensors sharded on dim-0, dealing with every parameter individually. This strategy permits simpler manipulation of parameters, for instance per-weight studying price, communication-free sharded state dicts, and less complicated meta-device initialization. The transition from FSDP1 to FSDP2 displays a shift in direction of extra versatile and environment friendly parameter dealing with in distributed coaching, addressing limitations of the flat-parameter strategy whereas doubtlessly introducing new optimization alternatives.

torchtitan assist for torch.compile

torch.compile is a key function in PyTorch that considerably boosts mannequin efficiency with minimal code modifications. By way of its just-in-time (JIT) compilation, it analyzes and transforms PyTorch code into extra environment friendly kernels. torchtitan helps torch.compile, which delivers substantial speedups, particularly for giant fashions and sophisticated architectures, through the use of methods like operator fusion, reminiscence planning, and automated kernel choice. That is enabled by setting compile = true within the mannequin’s TOML configuration file.

torchtitan assist for FP8 linear operations

torchtitan supplies assist for FP8 (8-bit floating level) computation that considerably reduces reminiscence footprint and enhances efficiency in LLM coaching. FP8 has two codecs, E4M3 and E5M2, every optimized for various elements of coaching. E4M3 presents greater precision, making it very best for ahead propagation, whereas E5M2, with its bigger dynamic vary, is best suited to backpropagation. When working at a decrease precision, FP8 has no affect on mannequin accuracy, which we display by convergence comparisons of the Meta Llama 3 8B pre-training at 2,000 steps. FP8 assist on torchtitan is thru the torchao library, and we allow FP8 by setting enable_float8_linear = true within the mannequin’s TOML configuration file.

torchtitan assist for FP8 all-gather

This function permits environment friendly communication of FP8 tensors throughout a number of GPUs, considerably lowering community bandwidth in comparison with bfloat16 all-gather operations. FP8 all-gather performs float8 casting earlier than the all-gather operation, lowering the message dimension. Key to its effectivity is the mixed absolute most (AMAX) AllReduce, which calculates AMAX for all float8 parameters in a single operation after the optimizer step, avoiding a number of small all-reduces. Much like FP8 assist, this additionally has no affect on mannequin accuracy, which we display by convergence comparisons of the Meta Llama 3 8B pre-training.

Pre-training Meta Llama 3 8B with torchtitan on Amazon SageMaker

SageMaker coaching jobs provide a number of key benefits that improve the pre-training means of Meta Llama 3-like mannequin architectures with torchtitan. It supplies a totally managed setting that simplifies large-scale distributed coaching throughout a number of cases, which is essential for effectively pre-training LLMs. SageMaker helps customized containers, which permits seamless integration of the torchtitan library and its dependencies, so all obligatory elements are available.

The built-in distributed coaching capabilities of SageMaker streamline the setup of multi-GPU and multi-node jobs, lowering the complexity sometimes related to such configurations. Moreover, SageMaker integrates with TensorBoard, enabling real-time monitoring and visualization of coaching metrics and offering beneficial insights into the pre-training course of. With these options, researchers and practitioners can focus extra on mannequin improvement and optimization quite than infrastructure administration, in the end accelerating the iterative course of of making and refining customized LLMs.

Resolution overview

Within the following sections, we stroll you thru easy methods to put together a customized picture with the torchtitan library, then configure a coaching job estimator perform to launch a Meta Llama 3 8B mannequin pre-training with the c4 dataset (Colossal Clear Crawled Corpus) on SageMaker. The c4 dataset is a large-scale net textual content corpus that has been cleaned and filtered to take away low-quality content material. It’s often used for pre-training language fashions.

Stipulations

Earlier than you start, be sure you have the next necessities in place:

Construct the torchtitan customized picture

SageMaker BYOC (Deliver Your Personal Container) permits you to use customized Docker containers to coach and deploy ML fashions. Usually, SageMaker supplies built-in algorithms and preconfigured environments for common ML frameworks. Nonetheless, there could also be circumstances the place you’ve distinctive or proprietary algorithms, dependencies, or particular necessities that aren’t obtainable within the built-in choices, necessitating customized containers. On this case, we have to use the nightly variations of torch, torchdata, and the torchao package deal to coach with FP8 precision.

We use the Amazon SageMaker Studio Picture Construct comfort package deal, which presents a command line interface (CLI) to simplify the method of constructing customized container pictures instantly from SageMaker Studio notebooks. This instrument eliminates the necessity for handbook setup of Docker construct environments, streamlining the workflow for knowledge scientists and builders. The CLI routinely manages the underlying AWS companies required for picture constructing, reminiscent of Amazon Easy Storage Service (Amazon S3), AWS CodeBuild, and Amazon Elastic Container Registry (Amazon ECR), permitting you to focus in your ML duties quite than infrastructure setup. It presents a easy command interface, handles packaging of Dockerfiles and container code, and supplies the ensuing picture URI to be used in SageMaker coaching and internet hosting.

Earlier than getting began, ensure your AWS Id and Entry Administration (IAM) execution position has the required IAM permissions and insurance policies to make use of the Picture Construct CLI. For extra info, see Utilizing the Amazon SageMaker Studio Picture Construct CLI to construct container pictures out of your Studio notebooks. We’ve got supplied the Jupyter pocket book to construct the customized container within the GitHub repo.

Full the next steps to construct the customized picture:

1. Set up the Picture Construct package deal with the next command:

! pip set up sagemaker-studio-image-build

2. To increase the pre-built picture, you should utilize the included deep studying libraries and settings with out having to create a picture from scratch:

FROM 763104351884.dkr.ecr.${REGION}.amazonaws.com/pytorch-training:2.3.0-gpu-py311-cu121-ubuntu20.04-sagemaker

3. Subsequent, specify the libraries to put in. You want the nightly variations of torch, torchdata, and the torchao libraries:

RUN pip set up --pre torch --force-reinstall --index-url https://obtain.pytorch.org/whl/nightly/cu121

RUN pip set up --pre torchdata --index-url https://obtain.pytorch.org/whl/nightly

#set up torchtitan dependencies
RUN pip set up --no-cache-dir 
datasets>=2.19.0 
tomli>=1.1.0 
tensorboard 
sentencepiece 
tiktoken 
blobfile 
tabulate

#set up torchao package deal for FP8 assist
RUN pip set up --pre torchao --index-url https://obtain.pytorch.org/whl/nightly/cu121
#Show put in packages for reference
RUN pip freeze

4. Use the Picture Construct CLI to construct and push the picture to Amazon ECR:

!sm-docker construct --repository torchtitan:newest . You’re now prepared to make use of this picture for pre-training fashions with torchtitan in SageMaker.

Put together your dataset (optionally available)

By default, the torchtitan library makes use of the allenai/c4 “en” dataset in its coaching configuration. That is streamed instantly throughout coaching utilizing the HuggingFaceDataset class. Nonetheless, you could need to pre-train the Meta Llama 3 fashions by yourself dataset residing in Amazon S3. For this function, we’ve ready a pattern Jupyter pocket book to obtain the allenai/c4 “en” dataset from the Hugging Face dataset hub to an S3 bucket. We use the SageMaker InputDataConfiguration to load the dataset to our coaching cases within the later part. You possibly can obtain the dataset with a SageMaker processing job obtainable within the pattern Jupyter pocket book.

Launch your coaching with torchtitan

Full the next steps to launch your coaching:

1. Import the mandatory SageMaker modules and retrieve your work setting particulars, reminiscent of AWS account ID and AWS Area. Make sure that to improve the SageMaker SDK to the newest model. This may require a SageMaker Studio kernel restart.

%pip set up --upgrade "sagemaker>=2.224"
%pip set up sagemaker-experiments

import os
import boto3
import sagemaker
from sagemaker import get_execution_role
from sagemaker.pytorch import PyTorch

position = get_execution_role()
print(f"SageMaker Execution Position: {position}")

consumer = boto3.consumer("sts")
account = consumer.get_caller_identity()["Account"]
print(f"AWS account: {account}")

session = boto3.session.Session()
area = session.region_name
print(f"AWS area: {area}")

sm_boto_client = boto3.consumer("sagemaker")
sagemaker_session = sagemaker.session.Session(boto_session=session)

default_bucket = sagemaker_session.default_bucket()
print("Default bucket for this session: ", default_bucket)

2. Clone the torchtitan repository and put together the coaching setting. Create a supply listing and transfer the mandatory dependencies from the torchtitan listing. This step makes positive you’ve all of the required recordsdata to your coaching course of.

git clone https://github.com/pytorch/torchtitan.git
mkdir torchtitan/src
!mv  torchtitan/torchtitan/ torchtitan/train_configs/ torchtitan/prepare.py  torchtitan/src/

3. Use the next command to obtain the Meta Llama 3 tokenizer, which is important for preprocessing your dataset. Present your Hugging Face token.

   python torchtitan/src/torchtitan/datasets/download_tokenizer.py --repo_id meta-llama/Meta-Llama-3-8B --tokenizer_path "authentic" --hf_token="YOUR_HF_TOKEN"

One of many key benefits of torchtitan is its easy configuration by means of TOML recordsdata. We modify the Meta Llama-3-8b TOML configuration file to allow monitoring and optimization options.

4. Allow TensorBoard profiling for higher insights into the coaching course of:

[metrics]
log_freq = 10
enable_tensorboard = true
save_tb_folder = "/choose/ml/output/tensorboard"

5. Allow torch.compile for improved efficiency:

6. Allow FP8 for extra environment friendly computations:

float8]
enable_float8_linear = true

7. Activate FP8 all-gather for optimized distributed coaching:

enable_fsdp_float8_all_gather= true
precompute_float8_dynamic_scale_for_fsdp = true

8. To watch the coaching progress, arrange TensorBoard output. This lets you visualize the coaching metrics in actual time, offering beneficial insights into how the mannequin is studying.

from sagemaker.debugger import TensorBoardOutputConfig

LOG_DIR="/choose/ml/output/tensorboard"
tensorboard_output_config = TensorBoardOutputConfig(
s3_output_path=f"s3://sagemaker-{area}-{account}/tensorboard/",
container_local_output_path=LOG_DIR
)

9. Arrange the info channels for SageMaker coaching. Create TrainingInput objects that time to the preprocessed dataset in Amazon S3, so your mannequin has entry to the coaching knowledge it wants.

#replace the trail under the s3 dataset path from operating the earlier Jupyter Pocket book from Step 2
training_dataset_location = "<PATH-TO-DATASET>" 

s3_train_bucket = training_dataset_location

if s3_train_bucket != None:
   prepare = sagemaker.inputs.TrainingInput(s3_train_bucket, distribution="FullyReplicated", s3_data_type="S3Prefix")
   data_channels = {"prepare": prepare}

10. With all of the items in place, you’re able to create the SageMaker PyTorch estimator. This estimator encapsulates all of the configurations, together with the customized container, hyperparameters, and useful resource allocations.

import os

from time import gmtime, strftime

hyperparameters = {
   "config_file": "train_configs/llama3_8b.toml"
}
timestamp = strftime("%Y-%m-%d-%H-%M", gmtime())


estimator = PyTorch(
   base_job_name=f'llama3-8b-{timestamp}',
   entry_point="prepare.py",
   image_uri="<PATH-TO-IMAGE-URI>",
   source_dir=os.path.be a part of(os.getcwd(), "src"),
   position=position,
   instance_type="ml.p5.48xlarge",
   volume_size=800,
   instance_count=4,
   hyperparameters=hyperparameters,
   use_spot_instances = False,
   sagemaker_session=sagemaker_session,
   tensorboard_output_config=tensorboard_output_config,
   distribution={
   'torch_distributed': {'enabled': True},
   },
  
)

11. Provoke the mannequin coaching on SageMaker:

estimator.match(inputs=data_channels)

Efficiency numbers

The next desk summarizes the efficiency numbers for the assorted coaching runs with totally different optimizations.

The efficiency outcomes present clear optimization progress in Meta Llama 3 8B pre-training. torch.compile() delivered an 10.67% speedup, and FP8 linear operations tripled this to 33%. Including FP8 all-gather additional elevated the speedup to 38.23% over the baseline. This development demonstrates how combining optimization methods considerably enhances coaching effectivity.

The next determine illustrates the stepwise efficiency beneficial properties for Meta Llama 3 8B pre-training on torchtitan with the optimizations.

These optimizations didn’t have an effect on the mannequin’s coaching high quality. The loss curves for all optimization ranges, together with the baseline, torch.compile(), FP8 linear, and FP8 all-gather configurations, remained constant all through the coaching course of, as proven within the following determine.

The next desk showcases the constant loss worth with the totally different configurations.

Clear up

After you full your coaching experiments, clear up your assets to keep away from pointless prices. You can begin by deleting any unused SageMaker Studio assets. Subsequent, take away the customized container picture from Amazon ECR by deleting the repository you created. When you ran the optionally available step to make use of your personal dataset, delete the S3 bucket the place this knowledge was saved.

Conclusion

On this put up, we demonstrated easy methods to effectively pre-train Meta Llama 3 fashions utilizing the torchtitan library on SageMaker. With torchtitan’s superior optimizations, together with torch.compile, FP8 linear operations, and FP8 all-gather, we achieved a 38.23% acceleration in Meta Llama 3 8B pre-training with out compromising the mannequin’s accuracy.

SageMaker simplified the large-scale coaching by providing seamless integration with customized containers, easy scaling throughout a number of cases, built-in assist for distributed coaching, and integration with TensorBoard for real-time monitoring.

Pre-training is an important step in creating highly effective and adaptable LLMs that may successfully sort out a variety of duties and functions. By combining the newest PyTorch distributed coaching options in torchtitan with the scalability and suppleness of SageMaker, organizations can use their proprietary knowledge and area experience to create sturdy and high-performance AI fashions. Get began by visiting the GitHub repository for the full code instance and optimize your LLM pre-training workflow.

Particular thanks

Particular due to Gokul Nadathur (Engineering Supervisor at Meta), Gal Oshri (Principal Product Supervisor Technical at AWS) and Janosch Woschitz (Sr. ML Resolution Architect at AWS) for his or her assist to the launch of this put up.

In regards to the Authors

Roy Allela is a Senior AI/ML Specialist Options Architect at AWS.He helps AWS prospects—from small startups to massive enterprises—prepare and deploy basis fashions effectively on AWS. He is obsessed with computational optimization issues and bettering the efficiency of AI workloads.

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

Trevor Harvey is a Principal Specialist in Generative AI at Amazon Internet 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 prospects to design and implement machine studying options and leads go-to-market methods for generative AI companies.

Much less Wright is an AI/Accomplice Engineer in PyTorch. He works on Triton/CUDA kernels (Accelerating Dequant with SplitK work decomposition); paged, streaming, and quantized optimizers; and PyTorch Distributed (PyTorch FSDP).

Wei Feng is a Software program Engineer on the PyTorch distributed crew. He has labored on float8 all-gather for FSDP2, TP (Tensor Parallel) in TorchTitan, and 4-bit quantization for distributed QLoRA in TorchTune. He’s additionally a core maintainer of FSDP2.