Enhancing AI Workloads with Amazon SageMaker HyperPod Task Governance

Accelerate Generative AI Innovation Through Topology-Aware Scheduling

Optimizing AI Workloads with SageMaker HyperPod Task Governance

In the ever-evolving landscape of artificial intelligence, optimizing resource allocation is crucial for driving innovation and reducing time to market. Today, we’re excited to announce a powerful new capability within Amazon SageMaker: HyperPod task governance. This feature provides a streamlined approach to enhance the training efficiency and minimize network latency for your AI workloads, particularly when deployed on Amazon Elastic Kubernetes Service (EKS) clusters.

What is SageMaker HyperPod Task Governance?

SageMaker HyperPod task governance simplifies how administrators manage accelerated compute allocations across teams and projects. With this enhanced capability, organizations can enforce task priority policies, ensuring efficient resource utilization. This allows data scientists to focus more on accelerating generative AI innovation rather than managing complex resource allocations.

The Importance of Network Configuration

Generative AI workloads typically require extensive network communication across Amazon Elastic Compute Cloud (EC2) instances. The network’s physical arrangement impacts both processing latency and workload runtime. For instance, if instances are within the same organizational unit—like a network node—they experience faster processing times compared to instances that are spread across different units. By minimizing network hops, organizations can significantly lower communication latency.

Introducing Topology-Aware Scheduling

One of the standout features of SageMaker HyperPod task governance is topology-aware scheduling. This allows users to consider the physical and logical arrangement of resources during job submissions, optimizing placement and enhancing communication efficiency. Key benefits of this approach include:

• Reduced Latency: By minimizing network hops, communication between instances is expedited.
• Improved Training Efficiency: This optimization leads to increased throughput and faster job completions.

How to Leverage Topology-Aware Scheduling

To effectively implement topology-aware scheduling, data scientists must first gain visibility into the topology information of all nodes in their cluster. This involves running scripts that display which instances reside on common network nodes, thereby allowing for informed decision-making regarding job submissions.

Setting Up Your Environment

To start with topology-aware scheduling, ensure you have the following prerequisites:

• An Amazon EKS cluster.
• A SageMaker HyperPod cluster with instances enabled for topology information.
• The SageMaker HyperPod task governance add-on (version 1.2.2 or later) installed.
• kubectl installed.
• (Optional) SageMaker HyperPod CLI installed.

Getting Node Topology Information

You can retrieve the node labels and network topology information for cluster instances using the following kubectl commands:

kubectl get nodes -L topology.k8s.aws/network-node-layer-1
kubectl get nodes -L topology.k8s.aws/network-node-layer-2
kubectl get nodes -L topology.k8s.aws/network-node-layer-3

This will provide insight into the layer structure of your cluster, allowing you to visualize the proximity of different instances.

Submitting Topology-Aware Tasks

Once you have determined the network node placements, you can submit tasks in two primary ways:

1. Modifying Your Kubernetes Manifest File

You can incorporate annotations into your existing manifest file to dictate pod placement. Here’s an example configuration:

apiVersion: batch/v1
kind: Job
metadata:
  name: test-task-job
spec:
  template:
    metadata:
      annotations:
        kueue.x-k8s.io/podset-required-topology: "topology.k8s.aws/network-node-layer-3"
    spec:
      containers:
        - name: dummy-job
          image: public.ecr.aws/docker/library/alpine:latest
          command: ["sleep", "3600s"]

2. Using the SageMaker HyperPod CLI

Alternatively, you can utilize the SageMaker HyperPod CLI for job submissions. Ensure you have the latest version installed, and use commands like:

hyp create hyp-pytorch-job \
--job-name test-pytorch-job-cli \
--image XXXXXXXXXXXX.dkr.ecr.us-west-2.amazonaws.com/ptjob:mnist \
--preferred-topology topology.k8s.aws/network-node-layer-3

Conclusion

As large language models and other AI workloads become more prevalent, the demand for efficient communication and data sharing across instances has never been higher. SageMaker HyperPod task governance combined with topology-aware scheduling offers a robust solution to meet these challenges.

We encourage you to explore this feature and integrate it into your AI training processes. Share your experiences and feedback in the comments below, as we continue to help organizations harness the power of generative AI.

About the Authors

This post was written by a talented team at AWS, including specialists in AI/ML technology, solutions architecture, and product management. Our collective goal is to empower organizations with cutting-edge AI capabilities, fostering innovation at every step of the journey.

Embrace the future of AI workloads with Amazon SageMaker HyperPod task governance and watch your innovations flourish!