Effectively practice fashions with massive sequence lengths utilizing Amazon SageMaker mannequin parallel

Massive language fashions (LLMs) have witnessed an unprecedented surge in recognition, with clients more and more utilizing publicly obtainable fashions equivalent to Llama, Steady Diffusion, and Mistral. Throughout various industries—together with healthcare, finance, and advertising—organizations are actually engaged in pre-training and fine-tuning these more and more bigger LLMs, which frequently boast billions of parameters and bigger enter sequence size. Though these developments supply outstanding capabilities, additionally they current important challenges. Longer sequence lengths and the sheer variety of trainable parameters demand modern approaches to mannequin growth and deployment. To maximise efficiency and optimize coaching, organizations regularly must make use of superior distributed coaching methods.

On this publish, we show how the Amazon SageMaker mannequin parallel library (SMP) addresses this want via help for brand new options equivalent to 8-bit floating level (FP8) mixed-precision coaching for accelerated coaching efficiency and context parallelism for processing massive enter sequence lengths, increasing the checklist of its present options.

We information you thru a step-by-step implementation, demonstrating speed up workloads with FP8 and work with longer sequence lengths utilizing context parallelism, with minimal code adjustments to your present coaching workflow.

The implementation of those new SMP options guarantees a number of benefits for purchasers working with LLMs. First, it could result in decrease prices to convergence, permitting for extra environment friendly use of sources throughout the coaching course of. This leads to diminished time to market, permitting organizations to deploy their optimized fashions extra shortly and acquire a aggressive edge. Second, it permits coaching with bigger dataset information, increasing the scope and complexity of duties that may be tackled.

The next sections take a deeper look into this.

Enterprise problem

Companies at present face a big problem when coaching LLMs effectively and cost-effectively. As fashions develop bigger and extra complicated, organizations are utilizing fine-tuning and steady pre-training methods to coach these fashions with domain-specific information, utilizing bigger sequence lengths that may vary from 8K to 128K tokens. These longer sequence lengths enable fashions to raised perceive long-range dependencies in textual content, generate extra globally coherent outputs, and deal with duties requiring evaluation of prolonged paperwork.

Though there exist varied methods equivalent to Absolutely Shared Information Parallelism (FSDP), tensor parallelism (TP), and pipeline parallelism to successfully practice fashions with billions of parameters, these strategies are primarily designed to distribute mannequin parameters, gradients, and optimizer states throughout GPUs, and so they don’t give attention to enter information–associated optimizations. This strategy reduces reminiscence strain and permits environment friendly coaching of huge fashions. Nevertheless, none of those methods successfully tackle partitioning alongside the sequence dimension. Consequently, coaching with longer sequence lengths can nonetheless result in out-of-memory (OOM) errors, regardless of utilizing FSDP.

Consequently, working with bigger sequence size would possibly end in reminiscence strain, and it typically requires modern approaches equivalent to FP8 and context parallelism.

How does SMP context parallelism and FP8 assist speed up mannequin coaching?

SMP addresses the challenges of reminiscence strain by offering an implementation of context parallelism, which is a parallelization method that partitions on the dimension of sequence size. Moreover, it could work along with different parallelism methods equivalent to FSDP and TP. SMP additionally implements FP8 for supported fashions equivalent to Llama. FP8 is a reduced-precision floating-point format that reinforces effectivity by enabling quicker matrix multiplications with out important accuracy loss. You need to use these methods collectively to coach complicated fashions which are orders of magnitude quicker and quickly iterate and deploy modern AI options that drive enterprise worth.

The next sections dive deep into the implementation particulars for every of those options in SMP.

Context parallelism

Context parallelism is a mannequin parallelism method to permit the mannequin to coach with lengthy sequences. It’s a parallelization scheme that partitions a mannequin’s activations alongside the sequence dimension. Throughout coaching with SMP context parallel technique, the inputs are partitioned alongside the sequence dimension earlier than being fed to the mannequin. With activations being partitioned alongside the sequence dimension, we have to contemplate how our mannequin’s computations are affected. For layers that don’t have inter-token dependency throughout computation, we don’t require particular concerns. In a transformer structure, such layers are the embedding layers and the multilayer perceptron (MLP) layers. The layers which have inter-token dependency are the eye layers. For the eye layer, as we see from the eye computation, Question projections (Q) must work together with the tokens of key (Okay) and worth (V) projections.

As a result of we solely have a partition of Okay and V, we require an AllGather operation to gather the keys and queries from different ranks. As detailed within the following determine, we contemplate a context parallel scheme with context parallel diploma 2 for a causal language mannequin. Thus GPU 0 has the primary half of the enter sequence and GPU 1 has the opposite half. Throughout ahead, the non-attention layers compute their activations as regular. For consideration computation, an AllGather operation is carried out for Okay and V throughout the context parallel ranks belonging to GPU 0 and GPU 1. To preserve reminiscence, the Okay and V tensors obtained from the AllGather operation are discarded after the eye computation is accomplished. Consequently, throughout the backward go, we require the identical AllGather operation for Okay and V. Moreover, after the eye backward go, a ReduceScatter operation is carried out to scatter the gradients to corresponding context parallel ranks.

In contrast to different mannequin parallel schemes equivalent to tensor parallelism, context parallelism retains the mannequin parameters intact. Thus, there are not any further communication collectives for parameters required for context parallelism.

Supported fashions

SMP helps context parallelism utilizing NVIDIA Transformer Engine, and it seamlessly integrates with different mannequin parallelism methods Absolutely Sharded Information Parallel and Tensor Parallelism. SMP v2.6 helps the Llama 3.1 (and prior Llama fashions) and Mistral mannequin architectures for context parallelism.

Blended Precision Coaching with FP8

As proven in determine under, FP8 is a datatype supported by NVIDIA’s H100 and H200 GPUs, permits environment friendly deep studying workloads. The FP8 format occupies solely 8 bits of reminiscence, half that of its BF16 or FP16 counterparts, considerably lowering computational prices for operations equivalent to matrix multiplication. The compute throughput for working matrix operations equivalent to multipliers and convolutions is considerably increased on 8-bit float tensors in comparison with 32-bit float tensors. FP8 precision reduces the info footprint and computational necessities, making it very best for large-scale fashions the place reminiscence and velocity are vital.

Delving deeper into FP8’s structure, we uncover two distinct subtypes: E4M3 and E5M2. The E4M3 configuration, with its 1 signal bit, 4 exponent bits, and three mantissa bits, presents superior precision however a restricted dynamic vary. This makes it very best for the ahead go in mannequin coaching. Conversely, E5M2, that includes 1 signal bit, 5 exponent bits, and a pair of mantissa bits, boasts a broader dynamic vary on the expense of diminished precision. This configuration excels within the backward go, the place precision is much less vital, however a wider vary proves advantageous.

The transition to combined precision coaching with FP16 or BF16 has traditionally necessitated static or dynamic loss-scaling to deal with convergence points that stemmed from diminished precision in gradient stream. This problem is additional amplified in FP8 attributable to its narrower vary. To fight this, the Transformer Engine launched an modern answer known as DelayedScaling. This system selects scaling components based mostly on the utmost noticed worth for every tensor from earlier iterations. Though DelayedScaling maximizes the efficiency advantages of FP8 computation, it does include a reminiscence overhead for storing the tensors’ most worth historical past. Nevertheless, regardless of the extra overhead, the improved throughput noticed with 8-bit tensor computations make this strategy helpful.

Supported fashions

SMP helps FP8 combined precision coaching utilizing NVIDIA Transformer Engine and retains compatibility with PyTorch MixedPrecision. Which means that you should utilize FP8 coaching for supported layers and half-precision utilizing PyTorch Computerized Blended Precision for others. SMP v2.6 helps the next mannequin architectures for FP8 coaching: Llama 3.1 (and prior Llama fashions), Mixtral, and Mistral.

Extra particulars about FP8 could be discovered at FP8 Codecs For Deep Studying.

Resolution overview

We will use SMP with each Amazon SageMaker Mannequin coaching jobs  and Amazon SageMaker HyperPod.

For this publish, we show SMP implementation on SageMaker trainings jobs.

Launching a machine studying (ML) coaching cluster with Amazon SageMaker coaching jobs is a seamless course of that begins with a simple API name, AWS Command Line Interface (AWS CLI) command, or AWS SDK interplay. After they’re initiated, SageMaker coaching jobs spin up the cluster, provisioning the desired quantity and kind of compute situations.

In our instance, we use a single ml.p5.48xlarge occasion, although we’re illustrating the usage of 4 GPUs for demonstration functions. The coaching information, securely saved in Amazon Easy Storage Service (Amazon S3), is copied to the cluster. Every report sequence (Seq0) is strategically break up into a number of subsequences and assigned to every GPU in our cluster.

Our implementation makes use of the FP8 capabilities of SMP to execute mannequin coaching on Nvidia H100 GPUs and showcases context parallelism capabilities. Due to the flexibleness of SageMaker, you may scale your compute sources as wanted, accommodating workloads throughout of a spread of sizes. SageMaker creates a resilient coaching cluster, handles orchestration, intently displays the infrastructure, and recovers from faults, offering a clean and uninterrupted coaching expertise. Moreover, the SageMaker coaching jobs cost-effective design routinely terminates the cluster upon completion of the coaching job, with billing calculated right down to the second of precise coaching time used. This mix of energy, flexibility, and cost-efficiency makes SageMaker an excellent service for ML practitioners of all ranges.

The next diagram reveals the answer structure.

The next walkthrough reveals you how one can practice a Llama 3.1 8B Instruct mannequin utilizing the PubMed tokenized dataset with a sequence size of roughly 16K tokens. We use SMP context parallelism implementation to allow coaching for this huge sequence size. We evaluate two approaches: one with out context parallelism and one other one with it. This comparability highlights the significance of context parallelism when working with LLMs and datasets containing lengthy sequences.

Moreover, we conduct a comparative run on p5.48xlarge situations with context parallelism enabled, each with FP8 enabled and disabled. This demonstration will showcase the incremental throughput advantages we are able to obtain by enabling FP8-based coaching alongside context parallelism.

In abstract, the implementation follows these 4 steps:

1. Arrange libraries and course of information
2. Run coaching with out context parallelism
3. Run coaching with context parallelism enabled to trace reminiscence optimizations
4. Run coaching with FP8 enabled to achieve additional efficiency

The next stream diagram reveals these 4 steps.

Stipulations

To carry out the answer, it’s worthwhile to have the next conditions in place:

1. Create a Hugging Face Consumer Entry Token and get entry to the gated repository meta-llama/Llama-3.1-8B on Hugging Face.
2. Request a Service Quota for 1x p4d.24xlarge and 1x ml.p5.48xlarge on Amazon SageMaker. To request a service quota improve, on the AWS Service Quotas console, select AWS providers, Amazon SageMaker, after which select one ml.p4d.24xlarge and one ml.p5.48xlarge coaching job utilization.
3. Create an AWS Id and Entry Administration (IAM) function with managed insurance policies AmazonSageMakerFullAccess, AmazonEC2FullAccess to offer required entry to SageMaker to run the examples.

This walkthrough is for demonstration functions solely. It is best to regulate this to your particular safety necessities for manufacturing. Adhere to the precept of least privilege whereas defining IAM insurance policies in manufacturing.

4. Create an Amazon SageMaker Studio area (discuss with Fast setup to Amazon SageMaker) to entry Jupyter notebooks.

Resolution walkthrough

To carry out the answer, use the directions within the following steps.

Arrange libraries and course of information

To arrange libraries and course of information, comply with these directions. The next stream diagram reveals step 1 highlighted.

1. Enter the next command to put in the related HuggingFace and SageMaker libraries:

   %pip set up --upgrade "sagemaker>=2.233"
   %pip set up "datasets==2.14.5"
   %pip set up transformers

2. Load the PubMed dataset and tokenize it

On this instance, we use the PubMed Scientific Papers dataset, containing 133,215 biomedical analysis articles. For our experiment, we choose 1,000 papers break up 80/20 for coaching and validation. Utilizing the Meta-LlaMA-3 tokenizer, we course of every paper into sequences of 16,384 tokens.

The dataset undergoes two principal processing steps: tokenization with Llama’s tokenizer and grouping into fixed-length chunks of 16,384 tokens utilizing utility perform group_texts. This uniform sequence size permits even distribution throughout GPUs whereas sustaining the pure construction of the scientific papers.

import datasets
from datasets import load_dataset, DatasetDict

# Load the PubMed dataset
pubmed_dataset = load_dataset(
    "scientific_papers",
    "pubmed",
    cache_dir="/dwelling/ec2-user/SageMaker/datasets",
    download_mode="force_redownload"
)

# Create a smaller subset of the dataset for our experiment
train_test = pubmed_dataset['train'].shuffle(seed=42).choose(vary(1000)).train_test_split(
    test_size=0.2,
    seed=42
)

lm_datasets = tokenized_datasets.map(
    group_texts,
    batched=True,
    desc=f"Grouping texts in chunks of {block_size}",
)

3. Put together information for the coaching job

On this part, we put together the PubMed dataset for SageMaker coaching by managing information transfers to Amazon S3. Each coaching and validation splits are transformed to JSON format and uploaded to designated S3 buckets, with separate paths for enter information and output artifacts.

if lm_datasets["train"] will not be None:
    train_dataset = lm_datasets["train"]
    train_dataset.to_json("./coaching.json")
    training_dataset_location = f"s3://{default_bucket}/dataset/practice/"

if lm_datasets["validation"] will not be None:
    eval_dataset = lm_datasets["validation"]
    eval_dataset.to_json("./validation.json")
    validation_dataset_location = f"s3://{default_bucket}/dataset/validation/"

4. Arrange coaching hyper parameters

On this configuration, we outline hyperparameters for coaching Llama on PubMed, masking reminiscence optimizations, coaching parameters, mannequin structure settings, and efficiency tuning. Beginning with conservative settings (batch measurement=1, BF16 precision), we set up a baseline configuration that will likely be modified to check totally different optimization methods, notably for context parallelism experiments.

hyperparameters = {
    # Reminiscence and optimization settings
    "activation_checkpointing": 1,
    "auto_wrap_policy": "transformer_auto_wrap_policy",
    ...
    
    # Coaching settings
    "train_batch_size": 1,
    "val_batch_size": 1,
    ...
    
    # Mannequin configuration
    "vocab_size": 128256, # Vocab measurement from Llama 3.1 config file on Hugging Face
    "hf_pretrained_model_name_or_dir": model_id,
    
    ...
    
}

Run coaching with out context parallelism

To run coaching with out context parallelism, comply with these directions. The next stream diagram reveals step 2 highlighted.

On this setup, we configure a baseline coaching job by disabling context parallelism and FP8 options, whereas maximizing reminiscence utilization via FP32 precision and bigger batch sizes. Every GPU processes the complete 16,384 token sequence with out splitting, and memory-saving options are disabled to show the constraints and potential reminiscence constraints when working with out superior optimizations equivalent to context parallelism and FP8.

instance_type= "p4d.24xlarge"
instance_count= 1
hybrid_shard_degree= 8

hyperparameters.replace({
    "use_smp_implementation": 0,  # Disable SMP/CP. Solely FSDP is lively
    "train_batch_size": 1,        # Batch measurement
    "max_context_width": 16384,   # Full sequence size
    "clean_cache": 0,
    "bf16": 1,                    # Use bf16
    ...
})

smp_estimator = PyTorch(
    entry_point="practice.py",
    hyperparameters=hyperparameters,
    ...
    instance_type=instance_type,
    volume_size=400,
    instance_type=instance_count,
    distribution={
        "torch_distributed": {
            "enabled": True,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,  # Allow mannequin parallelism however with minimal parameters
                "parameters": {
                    "hybrid_shard_degree": hybrid_shard_degree,
                    "delayed_parameter_initialization": True
                }
            }
        }
    },
    
   ...
)

smp_estimator.match(inputs=data_channels)

The results of not utilizing context parallelism with a big context width (16,384) means that we’ll get a CUDA out-of-memory error:

AlgorithmError: ExecuteUserScriptError: ExitCode 1 ErrorMessage “[rank3]: torch.OutOfMemoryError: CUDA out of reminiscence. Tried to allocate 7.83 GiB. GPU 3 has a complete capability of 39.38 GiB of which 5.53 GiB is free. Together with non-PyTorch reminiscence, this course of has 0 bytes reminiscence in use.

Run coaching with context parallelism enabled to trace reminiscence optimizations

To run coaching with context parallelism enabled to trace reminiscence optimizations, comply with these directions. The next stream diagram reveals step 3 highlighted.

On this configuration, we allow context parallelism whereas protecting FP8 disabled. By setting context parallel diploma to eight, we distribute the 16,384 token sequence throughout all obtainable GPUs for environment friendly processing. The setup consists of important context parallelism parameters and launches the coaching job in a background thread, permitting for unblocked pocket book execution whereas sustaining clear job identification for comparability with different configurations.

instance_type= "p4d.24xlarge"
instance_count= 1
hybrid_shard_degree= 8
context_parallel_degree=8

smp_estimator = PyTorch(
    ...
    entry_point="practice.py",
    instance_type=instance_type,
    instance_count=instance_count,
    distribution={
        "torch_distributed": {
            "enabled": True,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "context_parallel_degree": context_parallel_degree,
                    "hybrid_shard_degree": hybrid_shard_degree,
                    "delayed_parameter_initialization": True,
                }
            }
        }
    },
    ...
)

smp_estimator.match(inputs=data_channels)

The results of utilizing context parallelism with such a big context width is that the job efficiently completes, as proven within the following screenshot.

We additionally enabled delayed parameter initialization and hybrid sharding capabilities from SMP for each previous configurations. Delayed parameter initialization permits initializing massive fashions on a meta system with out attaching information. This will resolve restricted GPU reminiscence points whenever you first load the mannequin. This strategy is especially helpful for coaching LLMs with tens of billions of parameters, the place even CPU reminiscence won’t be adequate for initialization. Hybrid sharding is a reminiscence saving method that shards parameters inside the hybrid shard diploma (HSD) group and replicates parameters throughout teams. The HSD controls sharding throughout GPUs and could be set to an integer from 0 to world_size. This leads to diminished communication quantity as a result of costly AllGathers and ReduceScatters are solely executed inside a node, which carry out higher for medium-sized fashions.

Run coaching with FP8 enabled to achieve additional efficiency

To run coaching with FP8 enabled to achieve additional reminiscence efficiency, comply with these directions. The next stream diagram reveals step 4 highlighted.

On this absolutely optimized configuration, we allow each context parallelism and FP8 coaching utilizing a NVIDIA P5 occasion (ml.p5.48xlarge). This setup combines sequence splitting throughout GPUs with FP8 precision coaching, making a extremely environment friendly coaching surroundings. Utilizing P5 situations supplies the required {hardware} help for FP8 computation, with the outcome that we are able to maximize the advantages of each memory-saving methods.

instance_type= "p5.48xlarge"
instance_count= 1
hybrid_shard_degree= 8
context_parallel_degree=8

hyperparameters.replace({
    "use_smp_implementation": 1,  # Allow SMP/CP
    "max_context_width": 16384,   # Full sequence size
    "fp8": 1,  # Allow FP8 flag
    "distributed_backend": "nccl"  # Add this line to explicitly use NCCL
    ...

})

smp_estimator = PyTorch(
    ...
    entry_point="practice.py",
    instance_type=instance_type,
    instance_count=instance_count,
    distribution={
        "torch_distributed": {
            "enabled": True,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "context_parallel_degree": context_parallel_degree,
                    "hybrid_shard_degree": hybrid_shard_degree,
                    "delayed_parameter_initialization": True,
                }
            }
        }
    },
   ...
)

smp_estimator.match(inputs=data_channels)

Begin coaching with context parallelism, with out FP8 (on a P5 occasion)

To do a good comparability with and with out FP8, we’ll do one other run with out FP8 however with context parallelism on a P5.48xlarge occasion and evaluate the throughputs for each runs.

instance_type= "p5.48xlarge"
instance_count= 1
hybrid_shard_degree= 8
context_parallel_degree=8

hyperparameters.replace({
    "use_smp_implementation": 1,  # Allow SMP/CP
    "max_context_width": 16384,   # Full sequence size
    "bf16": 1,                    # Use BF16
    "distributed_backend": "nccl"  # Add this line to explicitly use NCCL
    ...
})

# This stays the identical as within the earlier step
smp_estimator = PyTorch(
    ...
    )
    
smp_estimator.match(inputs=data_channels)

If we evaluate each runs, we are able to inform that the velocity of the identical context parallelism enabled job with FP8 is nearly 10 instances quicker

With FP8, velocity is round 14.6 samples/second, as proven within the following screenshot.

With out FP8, velocity is round 1.4 samples/second, as proven within the following screenshot.

The next desk depicts the throughput increment you get in every of the listed instances. All these instances are run on a P5.48xLarge.

The throughput could fluctuate based mostly on components such because the context width or batch measurement. The next numbers are what we now have noticed in our testing.

Cleanup

To scrub up your sources to keep away from incurring extra costs, comply with these steps:

1. Delete any unused SageMaker Studio sources.
2. Optionally, delete the SageMaker Studio area.
3. Delete any S3 buckets created
4. Confirm that your coaching job isn’t working anymore! To take action, in your SageMaker console, select Coaching and verify Coaching jobs.

To be taught extra about cleansing up your sources provisioned, try Clear up.

Conclusion

On this publish, we demonstrated the method of organising and working coaching jobs for the PubMed dataset utilizing the Llama 3.1 8B Instruct mannequin, each with and with out context parallelism. We additionally showcased allow FP8 based mostly coaching for even quicker throughputs.

Key takeaways:

• For datasets which have lengthy sequence lengths, we observe that utilizing context parallelism helps keep away from OOM errors.
• For quicker coaching, we are able to allow FP8 based mostly coaching and mix it with context parallelism to get elevated throughput instances. On this pocket book, we noticed that the throughput goes up tenfold if we allow FP8 with context parallelism.

As subsequent steps, check out the above instance by following the pocket book steps at sagemaker-distributed-training-workshop.

Particular because of Roy Allela, Senior AI/ML Specialist Options Architect for his help on the launch of this publish.

Concerning the Authors

Kanwaljit Khurmi is a Principal Worldwide Generative AI Options Architect at AWS. He collaborates with AWS product groups, engineering departments, and clients to offer steerage and technical help, serving to them improve the worth of their hybrid machine studying options on AWS. Kanwaljit focuses on aiding clients with containerized purposes and high-performance computing options.

Surya Kari is a Senior Generative AI Information Scientist at AWS. With a background in laptop imaginative and prescient and AI gadgets, his present specializations embody LLM coaching, multi-modal RAG, vision-language fashions, and edge computing.

Arun Kumar Lokanatha is a Senior ML Options Architect with the Amazon SageMaker staff. He focuses on LLM coaching workloads, serving to clients construct LLM workloads utilizing SageMaker HyperPod, SageMaker coaching jobs, and SageMaker distributed coaching. Exterior of labor, he enjoys working, climbing, and cooking.

Suhit Kodgule is a Software program Improvement Engineer with the AWS Synthetic Intelligence group engaged on deep studying frameworks. In his spare time, he enjoys climbing, touring, and cooking.

Anirudh Viswanathan is a Sr Product Supervisor, Technical – Exterior Providers with the SageMaker Coaching staff. He holds a Masters in Robotics from Carnegie Mellon College, an MBA from the Wharton Faculty of Enterprise, and is known as inventor on over 40 patents. He enjoys long-distance working, visiting artwork galleries, and Broadway reveals.