New efficiency enhancements in Amazon SageMaker mannequin parallel library

Basis fashions are massive deep studying fashions educated on an unlimited amount of knowledge at scale. They are often additional fine-tuned to carry out quite a lot of downstream duties and kind the core spine of enabling a number of AI purposes. Probably the most distinguished class is large-language fashions (LLM), together with auto-regressive fashions reminiscent of GPT variants educated to finish pure textual content. LLMs usually comprise billions of parameters, making them not often match on one single accelerator, and require mannequin parallelism strategies. One other class is diffusion fashions, notably Steady Diffusion, that has pushed AI picture era to an unprecedented milestone the place outstanding visuals may be generated from a easy textual content description. Diffusion fashions are usually a lot smaller than LLMs and distributed coaching stays to play a vital position in facilitating growth.

SageMaker mannequin parallel (SMP) library is a large-model coaching resolution out there on Amazon SageMaker platform. It may be built-in with PyTorch fashions to simply apply a variety of state-of-the-art large-model distributed coaching strategies to coach at scale. Earlier this yr, SMP launched sharded knowledge parallelism, a distributed coaching approach powered by Amazon in-house MiCS expertise beneath the hood. Sharded knowledge parallel shards mannequin parameters, gradients, and optimizer states throughout data-parallel staff. MiCS performs various optimizations together with scale-aware partitioning to supply near-linear scalability.  In Prepare gigantic fashions with near-linear scaling utilizing sharded knowledge parallelism, we shared that sharded knowledge parallel in SMP achieved  39.7% velocity up in comparison with DeepSpeed ZeRO-3 on a 30B parameter GPT-2 mannequin with sequence size 2048.

To assist our prospects additional reduce coaching prices and speed up time-to-market, we’re thrilled to introduce two new efficiency enhancements in SageMaker mannequin parallel — SMDDP Collectives and FlashAttention. SMDDP Collectives is essentially the most performant collective library on AWS infrastructure for giant mannequin coaching supplied by SageMaker distributed knowledge parallel library. FlashAttention is launched in Dao et al., which re-implements the eye mechanism in an IO-aware method, decreasing the reminiscence bandwidth requirement and saving on consideration velocity and reminiscence footprint. These two parts collectively push our sharded knowledge parallel approach to be 30.58% quicker when coaching a 100B parameter GPT-NeoX mannequin on 32 p4d.24xlarge cases. For purchasers who’re already utilizing sharded knowledge parallel on supported fashions, no code adjustments are needed to learn from the efficiency enhance supplied by these newest options. Stability AI, the inventor of the Steady Diffusion household of fashions that confirmed unparalleled picture era skills, selected to make use of SMP to construct basis fashions. With SMP,  Stability AI achieved 163 TFLOPs per GPU for a 13B-parameter GPT-NeoX on 32 p4d.24xlarge cases, a 58% velocity up in comparison with DeepSpeed. You’ll be able to be taught extra about Stability AI’s mission and partnership with AWS within the discuss of Stability AI CEO at AWS re:Invent 2022 or on this weblog publish.

“Our mission at Stability AI is to construct the inspiration to activate humanity’s potential via AI. To realize this mission, we have to effectively prepare open-source basis fashions on tons of of accelerated compute cases. We depend on SageMaker and its distributed coaching libraries to optimize efficiency and implement state-of-the-art methods to shard fashions and knowledge throughout our coaching cluster. These optimizations cut back our coaching prices, assist us meet buyer wants quicker, and velocity up the event of recent fashions.”

— Emad Mostaque, Founder and CEO of Stability AI.

On this weblog publish, we’ll first current our newest efficiency enhancements within the SageMaker mannequin parallel library. Then, we’ll revisit how one can prepare foundational fashions utilizing sharded knowledge parallel.  Lastly, we’ll benchmark efficiency of 13B, 50B, and 100B parameter auto-regressive fashions and wrap up with future work.

New efficiency enhancements in  SageMaker mannequin parallel library

Ranging from AWS Deep Studying Containers (DLC) PyTorch 1.12.1, SageMaker mannequin parallel library v1.13 comes with the next two new parts which might be vital in enhancing coaching efficiency. They’re at present out there on ml.p4d.24xlarge occasion with Elastic Cloth Adapter (EFA) enabled:

1. AWS-optimized AllGather from SMDDP Collectives

In sharded knowledge parallel, since solely a shard of the mannequin state is current on a GPU, an AllGather collective is required to collect the complete set of parameters from throughout all GPUs within the sharding group throughout ahead or backward cross computations. Within the earlier variations of SageMaker mannequin parallel, we used NVIDIA Collective Communications Library (NCCL) for these collectives. Nevertheless, NCCL is a basic goal collective communications library not designed for AWS infrastructure, which results in sub-optimal efficiency even with EFA enabled.

Beforehand, we had developed the SMDDP Collectives library that supplied an AWS-optimized implementation of the All-Cut back collective to speedup efficiency of pure knowledge parallel coaching. To enhance the efficiency of enormous mannequin coaching with sharded knowledge parallelism, we expanded the SMDDP Collectives library to incorporate an optimized implementation of the AllGather collective. The important thing benefit of SMDDP Collectives AllGather is that it adopts an all-to-all-type communication sample for inter-node communication, enabling our collective to have high-throughput and be much less latency-sensitive. As well as, our AllGather collective offloads the communication-related processing to the CPU, thereby releasing up beneficial GPU cycles for gradient computation, resulting in important efficiency enchancment particularly on massive fashions.

2. FlashAttention

In fashionable transformer structure, one of many largest sources of reminiscence consumption is the activation footprint within the self-attention layer. It is because every consideration head computes an SxS consideration matrix for every enter, the place S is the sequence size, and this matrix goes via a number of operations, reminiscent of dropout, softmax, and matrix multiplication, with every intermediate output requiring reminiscence area to be used in back-propagation.

FlashAttention (Dao et al.) is a current innovation from HazyResearch in Stanford that re-implements the self-attention mechanism in an I/O-aware method. The principle perception behind FlashAttention is that the self-attention mechanism is bottlenecked by reminiscence bandwidth to and from GPU excessive bandwidth reminiscence (HBM). Which means the self-attention layer may be computed in chunks throughout the sequence dimension, with every chunk going via your entire self-attention pipeline at a time. The intermediate outcomes for a piece are saved on the high-bandwidth SRAM, avoiding the costly round-trip to the HBM for each iteration. Though a naive implementation would run into the difficulty of the cross-chunk dependency on the softmax layer, FlashAttention introduces a intelligent implementation that side-steps this dependency. Mixed with re-computation in backward cross, FlashAttention ends in substantial reminiscence financial savings and efficiency enchancment (25% quicker coaching for GPT-NeoX 13B over 16 p4d nodes), as a consequence of avoidance of the HBM round-trip and storage of SxS matrices. Yow will discover visuals and extra explanations in HazyResearch’s FlashAttention repository.

Prepare basis fashions at scale with SageMaker mannequin parallel

To coach basis fashions with SMP powered by SMDDP Collectives, there’s no extra adjustments required in your sharded knowledge parallel coaching jobs. When you’re new to utilizing sharded knowledge parallel, observe this entire tutorial pocket book and weblog publish that can stroll you thru your entire course of, from knowledge processing, defining and submitting coaching jobs, to monitoring coaching logs. A ready-to-use coaching script for GPT-2 mannequin may be discovered at train_gpt_simple.py. For coaching a unique mannequin sort, you’ll be able to observe the API doc to find out about how one can apply SMP APIs.

We spotlight the important thing hyperparameters within the PyTorch Estimator of a sharded knowledge parallel coaching job as under. The hyperparameter ddp_dist_backend in smp_options now has a brand new choice, "auto" , as its default worth. With "auto", SMP will use AWS-optimized AllGather for sharded knowledge parallelism jobs and fall again to NCCL in any other case. You’ll be able to discuss with this doc for supported configurations. If you wish to run sharded knowledge parallel in SMP particularly with NCCL because the communication backend of selection, you’ll be able to set “ddp_dist_backend" to "nccl" in smp_options.

import sagemaker
from sagemaker.pytorch import PyTorch

smp_options = {
    "enabled": True,
    "parameters": {
        "ddp": True,
        "ddp_dist_backend": "auto", #OR "nccl" to disable SMDDP Collectives
        # To allow sharded knowledge parallelism.
        # Right here we shard mannequin states throughout 128 GPUs.
        "sharded_data_parallel_degree": 128,  
    }
}

smp_estimator = PyTorch(
    entry_point="train_gpt_simple.py",
    position=sagemaker.get_execution_role(),
    instance_type="ml.p4d.24xlarge",
    instance_count=32,
    distribution={
        "smdistributed": {"modelparallel": smp_options},
        ...
    },
    ...
)

smp_estimator.match(inputs=data_channels)

With the most recent SMPv1.13 launch, the sharded knowledge parallel coaching approach helps FlashAttention for in style fashions together with BERT, RoBERTa, GPT-2, GPT-J, GPT-Neo and GPT-NeoX out-of-the-box. That is enabled by passing tensor_parallelism=True throughout mannequin creation with out setting tensor_parallel_degree. Yow will discover an instance in the identical coaching script train_gpt_simple.py .

Benchmarking efficiency

We benchmarked sharded knowledge parallelism within the SageMaker mannequin parallel library on three completely different scales of fashions to grasp how the 2 new options, FlashAttention and AWS-optimized AllGather, contribute to efficiency enchancment. Placement group isn’t required to breed these benchmarks on SageMaker.

13B parameter GPT-NeoX

On this setting, we concentrate on understanding the efficiency acquire contributed by FlashAttention and we depart AWS-optimized AllGather out of the image. Utilizing flash consideration saves substantial GPU reminiscence, which helps us improve batch measurement or cut back sharding diploma, thereby enhancing efficiency. Because the under outcomes present, we noticed a median of about 20.4% speedup in SMP with flash consideration for 13B parameter GPT-NeoX mannequin on varied configurations throughout 16-64 p4d nodes. Reminiscence utilization throughout normal consideration computation scales in a quadratic method with a rise in sequence size, however FlashAttention has reminiscence utilization linear in sequence size. Therefore FlashAttention is much more useful as sequence size will increase and makes it attainable to make use of bigger sequence lengths. Being memory-efficient with out buying and selling off mannequin high quality, FlashAttention has gained traction rapidly within the massive mannequin coaching neighborhood up to now months together with integration with Hugging Face Diffusers and Mosaic ML.

50B parameter Bloom

Now, we have a look at how AWS-optimized AllGather from SMDDP Collectives speedup massive mannequin coaching with SMP. We benchmark a 50B-parameter Bloom mannequin and examine the efficiency with and with out AWS-optimized AllGather collective. We observe that SMDDP collectives hastens mannequin coaching by upto 40% throughout 32 nodes to 64 nodes coaching jobs. SMDDP collectives assist obtain higher efficiency as a consequence of higher utilization of the 400 Gbps community bandwidth out there with p4d.24xlarge cases. This coupled with the design selection to dump communication-related processing to the CPU, helps obtain good compute-to-network overlap resulting in optimized efficiency. Compute-to-network overlap particularly turns into vital in massive fashions for the reason that measurement of knowledge communicated throughout nodes scales linearly with a rise within the mannequin measurement.

100B parameter GPT-NeoX

Lastly, we benchmark SMP with each of the most recent options enabled. It reveals that this new launch of SMP v1.13 is 30% quicker than the earlier model on a 100B-parameter GPT-NeoX mannequin.

For future work, we’ll be engaged on supporting an AWS-optimized Cut back-Scatter in SMDDP Collectives. The Cut back-Scatter collective is vital in averaging and sharding gradients computed within the backward cross. We count on this to additional velocity up SMP library sooner or later releases.

Conclusion

On this publish, we talk about the 2 newest efficiency enhancements for sharded knowledge parallel approach in SageMaker mannequin parallel library. LLMs present nice promise in enhancing the standard and re-usability of ML fashions. AWS groups are working carefully with prospects to maintain decreasing their coaching prices and time-to-market. Yow will discover extra SageMaker mannequin parallel examples in Amazon SageMaker Examples GitHub repo or attend our subsequent distributed coaching workshops. If you’re curious about dashing up massive mannequin coaching, take a look at these options and tell us what you construct!

In regards to the authors

Arjun Balasubramanian is a Senior Software program Engineer at AWS targeted on constructing high-performance, {hardware} accelerated collective communication algorithms for distributed deep studying. He’s broadly curious about techniques for large-scale machine studying and networking. Outdoors of labor, he enjoys touring and taking part in varied sports activities.

Zhaoqi Zhu is a Software program Growth Engineer at AWS, specializing in distributed deep studying techniques and dealing on the SageMaker Distributed Knowledge Parallel library. Outdoors of labor, Zhaoqi is enthusiastic about soccer and hopes to not obtain any purple card within the upcoming season.

Can Karakus is a Senior Utilized Scientist at AWS, optimizing large-scale distributed deep studying on AWS. His analysis pursuits cowl deep studying, distributed optimization, distributed techniques, and knowledge principle. Outdoors of labor, he enjoys biking, touring, studying and studying.

Rahul Huilgol is a Senior Software program Engineer at AWS. He works on distributed deep studying techniques, in direction of making it simple and performant to coach massive deep studying fashions within the cloud. In his spare time, he enjoys images, biking and gardening.

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

Fei Wu is a Software program Engineer at AWS. He works on distributed coaching for large-scale deep studying fashions on cloud. Outdoors of labor, he enjoys basketball, gaming and cooking.