Scaling distributed training with AWS Trainium and Amazon EKS

Current developments in deep studying have led to more and more giant fashions reminiscent of GPT-3, BLOOM, and OPT, a few of that are already in extra of 100 billion parameters. Though bigger fashions are usually extra highly effective, coaching such fashions requires important computational assets. Even with the usage of superior distributed coaching libraries like FSDP and DeepSpeed, it’s widespread for coaching jobs to require lots of of accelerator gadgets for a number of weeks or months at a time.

In late 2022, AWS introduced the overall availability of Amazon EC2 Trn1 cases powered by AWS Trainium—a purpose-built machine studying (ML) accelerator optimized to offer a high-performance, cost-effective, and massively scalable platform for coaching deep studying fashions within the cloud. Trn1 cases can be found in numerous sizes (see the next desk), with as much as 16 Trainium accelerators per occasion.

Trn1 cases can both be deployed as standalone cases for smaller coaching jobs, or in extremely scalable ultraclusters that assist distributed coaching throughout tens of 1000’s of Trainium accelerators. All Trn1 cases assist the standalone configuration, whereas Trn1 ultraclusters require trn1.32xlarge or trn1n.32xlarge cases. In an ultracluster, a number of Trn1 cases are co-located in a given AWS Availability Zone and are related with high-speed, low-latency, Elastic Cloth Adapter (EFA) networking that gives 800 Gbps of nonblocking community bandwidth per occasion for collective compute operations. The trn1n.32xlarge occasion sort, launching in early 2023, will improve this bandwidth to 1600 Gbps per occasion.

Many enterprise clients select to deploy their deep studying workloads utilizing Kubernetes—the de facto customary for container orchestration within the cloud. AWS clients usually deploy these workloads utilizing Amazon Elastic Kubernetes Service (Amazon EKS). Amazon EKS is a managed Kubernetes service that simplifies the creation, configuration, lifecycle, and monitoring of Kubernetes clusters whereas nonetheless providing the complete flexibility of upstream Kubernetes.

At this time, we’re excited to announce official assist for distributed coaching jobs utilizing Amazon EKS and EC2 Trn1 cases. With this announcement, now you can simply run large-scale containerized coaching jobs inside Amazon EKS whereas taking full benefit of the price-performance, scalability, and ease of use provided by Trn1 cases.

Together with this announcement, we’re additionally publishing an in depth tutorial that guides you thru the steps required to run a multi-instance distributed coaching job (BERT section 1 pre-training) utilizing Amazon EKS and Trn1 cases. On this put up, you’ll be taught concerning the resolution structure and assessment a number of key steps from the tutorial. Discuss with the official tutorial repository for the entire end-to-end workflow.

To observe alongside, a broad familiarity with core AWS companies reminiscent of Amazon Elastic Compute Cloud (Amazon EC2) and Amazon EKS is implied, and fundamental familiarity with deep studying and PyTorch could be useful.

Resolution structure

The next diagram illustrates the answer structure.

The answer consists of the next most important elements:

• An EKS cluster
• An EKS node group consisting of trn1.32xlarge cases
• The AWS Neuron SDK
• EKS plugins for Neuron and EFA
• An Amazon Elastic Container Registry (Amazon ECR) Rrepository
• A coaching container picture
• An Amazon FSx for Lustre file system
• A Volcano batch scheduler and etcd server
• The TorchX common job launcher
• The TorchX DDP module for Trainium

On the coronary heart of the answer is an EKS cluster that gives you with core Kubernetes administration performance by way of an EKS service endpoint. One of many advantages of Amazon EKS is that the service actively screens and scales the management airplane primarily based on load, which ensures excessive efficiency for giant workloads reminiscent of distributed coaching. Contained in the EKS cluster is a node group consisting of two or extra trn1.32xlarge Trainium-based cases residing in the identical Availability Zone.

The Neuron SDK is the software program stack that gives the motive force, compiler, runtime, framework integration (for instance, PyTorch Neuron), and consumer instruments that will let you entry the advantages of the Trainium accelerators. The Neuron gadget driver runs instantly on the EKS nodes (Trn1 cases) and gives entry to the Trainium chips from throughout the coaching containers which are launched on the nodes. Neuron and EFA plugins are put in throughout the EKS cluster to offer entry to the Trainium chips and EFA networking gadgets required for distributed coaching.

An ECR repository is used to retailer the coaching container photos. These photos comprise the Neuron SDK (excluding the Neuron driver, which runs instantly on the Trn1 cases), PyTorch coaching script, and required dependencies. When a coaching job is launched on the EKS cluster, the container photos are first pulled from Amazon ECR onto the EKS nodes, and the PyTorch employee containers are then instantiated from the pictures.

Shared storage is offered utilizing a high-performance FSx for Lustre file system that exists in the identical Availability Zone because the trn1.32xlarge cases. Creation and attachment of the FSx for Lustre file system to the EKS cluster is mediated by the Amazon FSx for Lustre CSI driver. On this resolution, the shared storage is used to retailer the coaching dataset and any logs or artifacts created in the course of the coaching course of.

The answer makes use of the TorchX universal job launcher to launch distributed coaching jobs inside Amazon EKS. TorchX has two vital dependencies: the Volcano batch scheduler and the etcd server. Volcano handles the scheduling and queuing of coaching jobs, whereas the etcd server is a key-value retailer utilized by TorchElastic for synchronization and peer discovery throughout job startup.

When a coaching job is launched utilizing TorchX, the launch command makes use of the offered TorchX distributed DDP module for Trainium to configure the general coaching job after which run the suitable torchrun instructions on every of the PyTorch employee pods. When a job is working, it may be monitored utilizing customary Kubernetes instruments (reminiscent of kubectl) or by way of customary ML toolsets reminiscent of TensorBoard.

Resolution overview

Let’s take a look at the vital steps of this resolution. All through this overview, we seek advice from the Launch a Multi-Node PyTorch Neuron Training Job on Trainium Using TorchX and EKS tutorial on GitHub.

Create an EKS cluster

To get began with distributed coaching jobs in Amazon EKS with Trn1 cases, you first create an EKS cluster as outlined within the tutorial on GitHub. Cluster creation might be achieved utilizing customary instruments reminiscent of eksctl and AWS CloudFormation.

Create an EKS node group

Subsequent, we have to create an EKS node group containing two or extra trn1.32xlarge cases in a supported Area. Within the tutorial, AWS CloudFormation is used to create a Trainium-specific EC2 launch template, which ensures that the Trn1 cases are launched with an acceptable Amazon Machine Picture (AMI) and the proper EFA community configuration wanted to assist distributed coaching. The AMI additionally consists of the Neuron gadget driver that gives assist for the Trainium accelerator chips. With the eksctl Amazon EKS administration device, you’ll be able to simply create a Trainium node group utilizing a fundamental YAML manifest that references the newly created launch template. For instance:

apiVersion: eksctl.io/v1alpha5
form: ClusterConfig

metadata:
  identify: my-trn1-cluster
  area: us-west-2
  model: "1.23"

iam:
  withOIDC: true

availabilityZones: ["us-west-xx","us-west-yy"]

managedNodeGroups:
  - identify: trn1-ng1
    launchTemplate:
      id: TRN1_LAUNCH_TEMPLATE_ID
    minSize: 2
    desiredCapacity: 2
    maxSize: 2
    availabilityZones: ["us-west-xx"]
    privateNetworking: true
    efaEnabled: true

Within the previous manifest, a number of attributes are configured to permit for the usage of Trn1 cases within the EKS cluster. First, metadata.area is ready to one of many Areas that helps Trn1 cases (presently us-east-1 and us-west-2). Subsequent, for availabilityZones, Amazon EKS requires that two Availability Zones be specified. One in all these Availability Zones should assist the usage of Trn1 cases, whereas the opposite might be chosen at random. The tutorial exhibits how one can determine which Availability Zones will allow for Trn1 instances within your AWS account. The identical Trn1-supporting Availability Zone should even be specified utilizing the availabiltyZones attribute related to the EKS node group. efaEnabled is ready to true to configure the nodes with the suitable EFA community configuration that’s required for distributed coaching. Lastly, the launchTemplate.id attribute related to the node group factors to the EC2 launch template created by way of AWS CloudFormation in an earlier step.

Assuming that you’ve got already utilized the CloudFormation template and put in the eksctl administration device, you’ll be able to create a Trainium-capable EKS node group by working the next code:

> eksctl create nodegroup -f TEMPLATE.yaml

Set up Kubernetes plugins for Trainium and EFA gadgets

With the node group in place, the following step is to put in Kubernetes plugins that present assist for the Trainium accelerators (by way of the Neuron plugin) and the EFA gadgets (by way of the EFA plugin). These plugins can simply be put in on the cluster utilizing the usual kubectl administration device as proven within the tutorial.

To make use of the TorchX common PyTorch launcher to launch distributed coaching jobs, two conditions are required: the Volcano batch scheduler, and the etcd server. Very like the Neuron and EFA plugins, we are able to use the kubectl device to put in Volcano and the etcd server on the EKS cluster.

Connect shared storage to the EKS cluster

Within the tutorial, FSx for Lustre is used to offer a high-performance shared file system that may be accessed by the varied EKS employee pods. This shared storage is used to host the coaching dataset, in addition to any artifacts and logs creating in the course of the coaching course of. The tutorial describes how one can create and fix the shared storage to the cluster utilizing the Amazon FSx for Lustre CSI driver.

Create a coaching container picture

Subsequent, we have to create a coaching container picture that features the PyTorch coaching script together with any dependencies. An instance Dockerfile is included within the tutorial, which includes the BERT pre-training script together with its software program dependencies. The Dockerfile is used to construct the coaching container picture, and the picture is then pushed to an ECR repository from which the PyTorch staff are capable of pull the picture when a coaching job is launched on the cluster.

Arrange the coaching information

Earlier than launching a coaching job, the coaching information is first copied to the shared storage quantity on FSx for Lustre. The tutorial outlines how one can create a short lived Kubernetes pod that has entry to the shared storage quantity, and exhibits how one can log in to the pod so as to obtain and extract the coaching dataset utilizing customary Linux shell instructions.

With the varied infrastructure and software program conditions in place, we are able to now give attention to the Trainium points of the answer.

Precompile your mannequin

The Neuron SDK helps PyTorch via an integration layer referred to as PyTorch Neuron. By default, PyTorch Neuron operates with just-in-time compilation, the place the varied neural community compute graphs inside a coaching job are compiled as they’re encountered in the course of the coaching course of. For bigger fashions, it may be extra handy to make use of the offered neuron_parallel_compile device to precompile and cache the varied compute graphs prematurely in order to keep away from graph compilation at coaching time. Earlier than launching the coaching job on the EKS cluster, the tutorial exhibits how one can first launch a precompilation job by way of TorchX utilizing the neuron_parallel_compile device. Upon completion of the precompilation job, the Neuron compiler may have recognized and compiled all the neural community compute graphs, and cached them to the shared storage quantity for later use in the course of the precise BERT pre-training job.

Launch the distributed coaching job

With precompilation full, TorchX is then used to launch a 64-worker distributed coaching job throughout two trn1.32xlarge cases, with 32 staff per occasion. We use 32 staff per occasion as a result of every trn1.32xlarge occasion incorporates 16 Trainium accelerators, with every accelerator offering 2 NeuronCores. Every NeuronCore might be accessed as a singular PyTorch XLA device within the coaching script. An instance TorchX launch command from the tutorial seems to be like the next code:

    torchx run 
    -s kubernetes --workspace="file:///$PWD/docker" 
    -cfg queue=check,image_repo=$ECR_REPO 
    lib/trn1_dist_ddp.py:generateAppDef 
    --name berttrain 
    --script_args "--batch_size 16 --grad_accum_usteps 32 
        --data_dir /information/bert_pretrain_wikicorpus_tokenized_hdf5_seqlen128 
        --output_dir /information/output" 
    --nnodes 2 
    --nproc_per_node 32 
    --image $ECR_REPO:bert_pretrain 
    --script dp_bert_large_hf_pretrain_hdf5.py 
    --bf16 True 
    --cacheset bert-large

The varied command line arguments within the previous TorchX command are described intimately within the tutorial. Nevertheless, the next arguments are most vital in configuring the coaching job:

• -cfg queue=check – Specifies the Volcano queue for use for the coaching job
• -cfg image_repo – Specifies the ECR repository for use for the TorchX container photos
• –script_args – Specifies any arguments that needs to be handed to the PyTorch coaching script
• –nnodes and –nproc_per_node – The variety of cases and staff per occasion to make use of for the coaching job
• –script – The identify of the PyTorch coaching script to launch throughout the coaching container
• –picture – The trail to the coaching container picture in Amazon ECR
• –bf16 – Whether or not or to not allow BF16 information sort

Monitor the coaching job

After the coaching job has been launched, there are numerous methods by which the job might be monitored. The tutorial exhibits how one can monitor fundamental coaching script metrics on the command line utilizing kubectl, how one can visually monitor coaching script progress in TensorBoard (see the next screenshot), and how one can monitor Trainium accelerator utilization utilizing the neuron-top device from the Neuron SDK.

Clear up or reuse the surroundings

When the coaching job is full, the cluster can then be reused or re-configured for extra coaching jobs. For instance, the EKS node group can shortly be scaled up utilizing the eksctl command so as to assist coaching jobs that require further Trn1 cases. Equally, the offered Dockerfile and TorchX launch instructions can simply be modified to assist further deep studying fashions and distributing coaching topologies.

If the cluster in now not required, the tutorial additionally consists of all steps required to take away the EKS infrastructure and associated assets.

Conclusion

On this put up, we explored how Trn1 cases and Amazon EKS present a managed platform for high-performance, cost-effective, and massively scalable distributed coaching of deep studying fashions. We additionally shared a complete tutorial exhibiting how one can run a real-world multi-instance distributed coaching job in Amazon EKS utilizing Trn1 cases, and highlighted a number of of the important thing steps and elements within the resolution. This tutorial content material can simply be tailored for different fashions and workloads, and gives you with a foundational resolution for distributed coaching of deep studying fashions in AWS.

To be taught extra about how one can get began with Trainium-powered Trn1 cases, seek advice from the Neuron documentation.

Concerning the Authors

Scott Perry is a Options Architect on the Annapurna ML accelerator staff at AWS. Based mostly in Canada, he helps clients deploy and optimize deep studying coaching and inference workloads utilizing AWS Inferentia and AWS Trainium. His pursuits embody giant language fashions, deep reinforcement studying, IoT, and genomics.

Lorea Arrizabalaga is a Options Architect aligned to the UK Public Sector, the place she helps clients design ML options with Amazon SageMaker. She can also be a part of the Technical Discipline Neighborhood devoted to {hardware} acceleration and helps with testing and benchmarking AWS Inferentia and AWS Trainium workloads.