Speed up large-scale AI coaching with Amazon SageMaker HyperPod coaching operator 

Giant-scale AI mannequin coaching faces important challenges with failure restoration and monitoring. Conventional coaching requires full job restarts when even a single coaching course of fails, leading to extra downtime and elevated prices. As coaching clusters develop, figuring out and resolving crucial points like stalled GPUs and numerical instabilities usually requires advanced customized monitoring code.

With Amazon SageMaker HyperPod you possibly can speed up AI mannequin growth throughout lots of or hundreds of GPUs with built-in resiliency, reducing mannequin coaching time by as much as 40%. The Amazon SageMaker HyperPod coaching operator additional enhances coaching resilience for Kubernetes workloads by way of pinpoint restoration and customizable monitoring capabilities.

On this weblog put up, we present you easy methods to deploy and handle machine studying coaching workloads utilizing the Amazon SageMaker HyperPod coaching operator, together with setup directions and a whole coaching instance.

Amazon SageMaker HyperPod coaching operator

The Amazon SageMaker HyperPod coaching operator helps you speed up generative AI mannequin growth by effectively managing distributed coaching throughout massive GPU clusters. The Amazon SageMaker HyperPod coaching operator makes use of built-in fault resiliency parts, comes packaged as an Amazon Elastic Kubernetes Service (Amazon EKS) add-on, and deploys the mandatory customized useful resource definitions (CRDs) to the HyperPod cluster.

Resolution overview

The next diagram depicts the structure of Amazon SageMaker HyperPod coaching operator.

The HyperPod coaching operator follows Kubernetes operator pattern and has the next main parts:

1. Customized Useful resource Definition (CRDs): HyperPodPyTorchJob defines the job specification (for instance, node rely, picture) and serves because the interface for patrons to submit jobs.
   apiVersion: sagemaker.amazonaws.com/v1
form: HyperPodPyTorchJob
2. RBAC insurance policies: Defines the actions the controller is allowed to carry out, corresponding to creating pods and managing HyperPodPyTorchJob sources.
3. Job controller: Listens to job creation and fulfills requests by creating job pods and pod managers.
4. Pod supervisor: Displays coaching course of well being on every pod. The variety of Pod Managers is decided by the variety of pods required by the job. One Pod Supervisor presently controls a number of hundred pods.
5. HyperPod elastic agent: Prospects set up the elastic agent into their coaching container. It orchestrates lifecycles of coaching employees on every container and communicates with the Amazon SageMaker HyperPod coaching operator. The HyperPod elastic agent is an extension of PyTorch’s ElasticAgent.

The job Controller makes use of fault detection parts such because the SageMaker HyperPod health-monitoring agent and node well being verify mechanisms like AWS retirement notices to replace job state and restore faults. It additionally depends on the HyperPod elastic agent to verify the standing of coaching processes for crashes and hung job detection.

When a HyperPodPyTorch job is submitted, the Amazon SageMaker HyperPod coaching operator spins up job pods together with pod supervisor pods that assist handle the coaching job lifecycle. The pod managers work together with the HyperPod elastic agent so that every one job pods preserve a wholesome state.

Advantages of utilizing the operator

The Amazon SageMaker HyperPod coaching operator may be put in as an EKS add-on in your cluster. The important thing advantages embrace:

• Centralized coaching course of monitoring and restart – The HyperPod coaching operator maintains a management aircraft with a worldwide view of well being throughout all ranks. When one rank encounters a difficulty, it broadcasts a cease sign to all ranks to forestall different ranks from failing individually at completely different occasions on account of collective communication timeout. This helps extra environment friendly fault detection and restoration.
• Centralized environment friendly rank project – A separate HyperPod rendezvous backend permits the HyperPod coaching operator to assign ranks instantly. This reduces initialization overhead by eliminating the necessity for worker-to-worker discovery.
• Unhealthy coaching node detection and job restart – The HyperPod coaching operator is totally built-in with the HyperPod EKS cluster resiliency options, serving to restart jobs or coaching processes on account of unhealthy nodes and {hardware} points in ML workloads. This reduces the necessity to self-manage job restoration options.
• Granular course of restoration – Slightly than restarting whole jobs when failures happen, the operator exactly targets and restarts solely coaching processes, lowering restoration occasions from tens of minutes to seconds. This makes HyperPod coaching operator job restoration time scale linearly as cluster measurement grows.
• Hanging job detection and efficiency degradation detection – Primarily based on coaching script log monitoring, the HyperPod coaching operator helps overcome problematic coaching eventualities together with stalled coaching batches, non-numeric loss values, and efficiency degradation by way of easy YAML configurations. For extra data see, Using the training operator to run jobs within the Amazon SageMaker AI Developer Information.

Coaching operator setup

This part walks by way of putting in the Amazon SageMaker HyperPod coaching operator as an Amazon EKS add-on.

Estimated Setup Time: 30-45 minutes

Conditions

Earlier than getting began, confirm that you’ve got the next sources and permissions.

Required AWS sources:

Required IAM permissions:

Required software program:

Set up directions

Earlier than working the set up steps under, you’ll have to first create a HyperPod cluster. For those who haven’t performed this one already please comply with the directions to create an EKS-orchestrated SageMaker HyperPod cluster to get began. Be sure to put in eks-pod-identity-agent add-on on the EKS cluster, by following the Set up the Amazon EKS Pod Identity Agent directions.

Set up cert-manager

First, set up the cert-manager add-on which is required for the HyperPod coaching operator:

1. Open the Amazon EKS console
2. Navigate to your EKS cluster and go to the Add-ons web page
3. On the Add-ons web page, find Get extra add-ons and navigate to the Neighborhood add-ons part
4. Discover the Cert Supervisor add-on, choose it, and select Subsequent
5. On the add-on configuration web page, proceed with default settings and select Subsequent
6. Preview all picks for the Cert Supervisor add-on and select Create
7. Anticipate the add-on standing to alter to Energetic earlier than continuing

Set up the HyperPod coaching operator add-on

As soon as cert-manager is energetic, set up the Amazon SageMaker HyperPod coaching operator:

1. Open the Amazon SageMaker console
2. Navigate to your cluster’s particulars web page
3. On the Dashboard tab, find Amazon SageMaker HyperPod coaching operator and select Set up

Throughout set up, SageMaker creates an IAM execution function with permissions just like the AmazonSageMakerHyperPodTrainingOperatorAccess managed coverage and creates a pod id affiliation between your Amazon EKS cluster and the brand new execution function.

Confirm set up

We now have now efficiently setup of the Amazon SageMaker HyperPod coaching operator. You may affirm that the pods are working by utilizing the next command:

kubectl -n aws-hyperpod get pods -l hp-training-control-plane=hp-training-operator-controller-manager

Your output ought to include the coaching operator controller as proven under:

NAME READY      STATUS         RESTARTS        AGE
hp-training-operator-hp-training-controller-manager-85c68bmd79b    1/1              Working         0                        24m

Arrange coaching job

Let’s run a PyTorch-based coaching instance on a Llama mannequin. We start by testing the next code base:

git clone 
cd awsome-distributed-training/tree/foremost/3.test_cases/pytorch/FSDP

These scripts present a simple option to get began with multinode FSDP coaching on EKS. It’s designed to be so simple as doable, requires no information preparation, and makes use of a container picture.

Subsequent, construct the docker container picture.

aws ecr-public get-login-password —area us-east-1 | docker login —username AWS —password-stdin public.ecr.aws/hpc-cloud
export REGION=$(aws ec2 describe-availability-zones —output textual content —question 'AvailabilityZones[0].[RegionName]')
export ACCOUNT=$(aws sts get-caller-identity —question Account —output textual content)
export REGISTRY=${ACCOUNT}.dkr.ecr.${REGION}.

docker construct -t ${REGISTRY}fsdp:pytorch2.5.1 .    

The above command works with linux based mostly environments, if you’re on a Mac, use buildx to focus on linux/amd64 structure:

docker buildx construct --platform linux/amd64 -t ${REGISTRY}fsdp:pytorch2.5.1 .

Push the picture to Amazon ECR:

# Create registry if wanted
REGISTRY_COUNT=$(aws ecr describe-repositories | grep "fsdp" | wc -l)
if [ "$REGISTRY_COUNT" -eq 0 ]; then
    aws ecr create-repository --repository-name fsdp
fi

# Login to registry
echo "Logging in to $REGISTRY ..."
aws ecr get-login-password | docker login --username AWS --password-stdin $REGISTRY

# Push picture to registry
docker picture push ${REGISTRY}fsdp:pytorch2.5.1

Be aware: Pushing the picture might take a while relying in your community bandwidth.

Information

For this instance, we’ll be utilizing the allenai/c4 dataset. As an alternative of downloading the entire thing, the create_streaming_dataloaders perform will stream the dataset from HuggingFace, so there’s no information prep required for working this coaching.

For those who’d prefer to as a substitute use your individual dataset, you are able to do so by formatting it as a HuggingFace dataset, and passing its location to the --dataset_path argument.

For the dataset, you have to a Hugging Face entry token. First, create a Hugging Face account. Then generate your access token with read permissions.

We’ll reference this token within the subsequent step by setting it as an atmosphere variable.

This instance makes use of envsubst to generate a Kubernetes manifest file from a template file and parameters. For those who don’t have envsubst in your growth atmosphere, set up it by following the installation instructions.

Launch Llama 3.1 8B coaching job

Subsequent, we generate the Kubernetes manifest and apply it to the cluster. Let’s navigate to the FSDP supply repo:

cd awsome-distributed-training/tree/foremost/3.test_cases/pytorch/FSDP/Kubernetes

Right here, we begin by creating atmosphere variables which can be utilized in our coaching job. Fill out the placeholders as per your cluster measurement.

cat << EOF > env_vars
export ACCOUNT_ID=<AWS_ACCOUNT_ID>
export REGION=<REGION>
export REGISTRY=${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com
export IMAGE_URI=${REGISTRY}/fsdp:pytorch2.5.1
export INSTANCE_TYPE=<INSTANCE TYPE> # ml.p5.48xlarge
export NUM_NODES=<NUMBER OF NODES> # 2
export GPU_PER_NODE=<NUMBER OF GPUS PER NODE> # 8
export EFA_PER_NODE=<NUMBER OF EFA PER NODE> # 32
export FI_PROVIDER=efa
export HF_TOKEN=<YOUR HF ACCESS TOKEN> # HF_xxxx
EOF

When you fill in env_vars after which supply variables:

You may apply yaml to submit the coaching job:

envsubst < llama3_1_8b-fsdp-hpto.yaml | kubectl apply -f -

You can too regulate the coaching parameters within the TRAINING_ARGS part of the llama3_1_8b-fsdp-hpto.yaml. Further parameters may be discovered beneath mannequin/arguments.py. Be aware that we use the identical listing for each --checkpoint_dir and --resume_from_checkpoint. If there are a number of checkpoints, --resume_from_checkpoint will robotically choose the newest one. This fashion if our coaching is interrupted for any motive, it is going to robotically decide up the newest checkpoint.

Moreover, you may also put together and submit your jobs suitable with the Amazon SageMaker HyperPod coaching operator by way of the HyperPod CLI and SDK capabilities which have been lately introduced, extra studying data on easy methods to use it’s obtainable in this development guide.

Monitor coaching job

To see the standing of your job, use the next command:

kubectl get hyperpodpytorchjobs

Use the next command to record the roles ran utilizing HyperPod coaching operator:

NAME              AGE
llama2-13b-fsdp   2m15s

kubectl get pods 

Use the next command to record all of the pods for the coaching jobs:

NAME                    READY  STATUS   RESTARTS AGE
llama2-13b-fsdp-pods-0  1/1    Working   0       13s
llama2-13b-fsdp-pods-1  1/1    Working   0       13s
llama2-13b-fsdp-pods-2  1/1    Working   0       13s
llama2-13b-fsdp-pods-3  1/1    Working   0       13s

To verify the pod logs run the under command to repeatedly stream the logs to stdout, use the next command:

kubectl logs -f llama2-13b-fsdp-pods-0

Configure log monitoring

With Amazon SageMaker HyperPod coaching operators customers can configure log patterns that the operator repeatedly screens. The HyperPod operator repeatedly appears to be like for the configured regex sample and stops the coaching job if it finds a violation. The llama3_1_13b-fsdp-hpto.yaml file that we used beforehand accommodates log monitoring configurations for monitoring Job begin hangs, grasp detection throughout coaching, and checkpoint creation failures as proven under:

  logMonitoringConfiguration:
      - identify: "JobStart"
        logPattern: ".*Loss:.*"
        expectedStartCutOffInSeconds: 240 # job ought to print loss inside 4 minutes of begin time
      - identify: "JobHangingDetection"
        logPattern: ".*Loss:.*"
        expectedRecurringFrequencyInSeconds: 300 # if subsequent batch is just not printed inside 300 seconds
      - identify: "NoSCheckpointingDetection"
        logPattern: ".*Accomplished checkpoint.*"
        expectedRecurringFrequencyInSeconds: 600 # If subsequent checkpoint add does not occur inside 10 minutes, mark it grasp.
        expectedStartCutOffInSeconds: 900 # Permit half-hour for first checkpoint add

And the corresponding code information in /src/practice.py have the mandatory log statements.

logger.data(
             "Batch %d Loss: %.5f, Velocity: %.2f samples/sec, lr: %.6f",  # pylint: disable=line-too-long
                    batch_idx,
                    loss_scalar,
                    throughput,
                    current_lr,
                )
                

Any time these metrics exhibit deviation from their anticipated values, the operator will detect it as a fault, and set off a restoration course of to re-execute the job, as much as a user-specified most variety of retries.

Moreover, the HyperPod coaching operator additionally helps integration with Amazon SageMaker Task Governance.

Integration with HyperPod Observability

SageMaker HyperPod affords a managed observability expertise by way of the newly launched the HyperPod Monitoring and Observability EKS add-on. The observability add-on robotically populates Kubeflow Coaching metrics in Grafana dashboards out of the field, however for HyperPod PyTorch job metrics, you would need to activate the superior coaching metrics which leverage the HyperPod coaching operator to indicate data round job downtime, job restoration and faults, and downtime.

To get these superior metrics, you possibly can discuss with Setting up the SageMaker HyperPod observability add-on. This helps to streamline the method of manually organising a scraper and constructing dashboards.

Clear up

To keep away from incurring pointless expenses, clear up the sources created on this walkthrough.

Delete coaching jobs

Take away all HyperPod coaching jobs:

kubectl delete hyperpodpytorchjobs --all

Confirm jobs are deleted:

kubectl get hyperpodpytorchjobs

Take away container photos

Delete the ECR repository and pictures:

aws ecr delete-repository --repository-name fsdp --force

Take away add-ons:

Take away the next add-ons:

Take away the Amazon SageMaker HyperPod coaching operator add-on:

1. Open the Amazon SageMaker console
2. Navigate to your cluster’s particulars web page
3. On the Add-ons tab, choose the Amazon SageMaker HyperPod coaching operator
4. Select Take away

Take away the cert supervisor add-on:

1. Open the Amazon EKS console
2. Navigate to your EKS cluster’s Add-ons web page
3. Choose Cert Supervisor and select Take away

Further clear up

Think about eradicating these sources if not wanted:

• Any persistent volumes created throughout coaching
• CloudWatch log teams (if you wish to retain logs, go away these)
• Customized IAM roles created particularly for this instance
• The HyperPod cluster itself (if not wanted).

Conclusion

As organizations proceed to push the boundaries of AI mannequin growth, instruments just like the Amazon SageMaker HyperPod coaching operator can be utilized to keep up effectivity and reliability at scale. Amazon SageMaker HyperPod coaching operator affords a strong answer to widespread challenges in massive mannequin coaching. Key takeaways embrace:

1. One-click set up by way of AWS SageMaker HyperPod cluster console user-interface.
2. Customized rendezvous backend eliminates initialization and employee synchronization overhead which ends up in quicker job begins and restoration.
3. Course of stage restarts maximize restoration effectivity when runtime faults happen.
4. Customizable grasp job detection throughout coaching.
5. Complete monitoring for early detection of coaching points.
6. Out-of-box integration with present HyperPod resiliency options.

To get began with the Amazon SageMaker HyperPod coaching operator, comply with the setup directions offered on this put up and discover the instance coaching job to know the way it can profit your particular use case. For extra data and greatest practices, go to the Amazon SageMaker documentation.

Concerning the authors

Arun Kumar Lokanatha is a Senior ML Options Architect with the Amazon SageMaker AI. He holds a Grasp’s diploma from UIUC with a specialization in Information science. He makes a speciality of Generative AI workloads, serving to clients construct and deploy LLM’s utilizing SageMaker HyperPod, SageMaker coaching jobs, and SageMaker distributed coaching. Outdoors of labor, he enjoys working, mountaineering, and cooking.

Haard Mehta is a Software program Engineer with Amazon’s SageMaker AI crew and holds a Grasp’s diploma in Laptop Science with a specialization in huge information methods from Arizona State College. He has intensive expertise constructing managed machine studying companies at scale, with a give attention to {hardware} resiliency and enabling clients to reach their AI use circumstances with out advanced infrastructure administration. Haard enjoys exploring new locations, pictures, cooking, and street journeys.

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