调度拓扑感知工作负载使用某中心SageMaker HyperPod任务治理

今天，我们很高兴宣布某中心SageMaker HyperPod任务治理的新功能，可帮助您优化AI工作负载的训练效率和网络延迟。SageMaker HyperPod任务治理简化了资源分配，并在某中心弹性Kubernetes服务集群上促进团队和项目间的高效计算资源利用。管理员可以管理加速计算分配并强制执行任务优先级策略，从而提高资源利用率。这有助于组织专注于加速生成式AI创新并缩短上市时间，而不是协调资源分配和重新规划任务。

生成式AI工作负载通常需要跨某中心弹性计算云实例进行广泛的网络通信，其中网络带宽会影响工作负载运行时间和处理延迟。这些通信的网络延迟取决于实例在数据中心分层基础设施中的物理位置。数据中心可以组织成嵌套的组织单元，例如网络节点和节点集，每个网络节点有多个实例，每个节点集有多个网络节点。例如，同一组织单元内的实例相比跨不同单元的实例具有更快的处理时间。这意味着实例间较少的网络跳数会带来更低的通信延迟。

为了通过考虑资源的物理和逻辑安排来优化生成式AI工作负载在SageMaker HyperPod集群中的放置，您可以在作业提交期间使用EC2网络拓扑信息。EC2实例的拓扑由一组节点描述，每个网络层有一个节点。网络拓扑标签提供以下关键优势：

• 通过最小化网络跳数并将流量路由到附近实例来减少延迟
• 通过优化跨网络资源的工作负载放置来提高训练效率

通过SageMaker HyperPod任务治理的拓扑感知调度，您可以使用拓扑网络标签来调度作业，优化网络通信，从而提高AI工作负载的任务效率和资源利用率。

在本文中，我们通过提交代表分层网络信息的作业，介绍使用SageMaker HyperPod任务治理进行拓扑感知调度。我们提供有关如何使用SageMaker HyperPod任务治理来优化作业效率的详细信息。

解决方案概述

数据科学家与SageMaker HyperPod集群交互。数据科学家负责在加速计算实例上进行模型的训练、微调和部署。确保数据科学家在与GPU集群交互时具有必要的容量和权限非常重要。

要实现拓扑感知调度，您首先确认集群中所有节点的拓扑信息，然后运行一个脚本，告诉您哪些实例位于相同的网络节点上，最后在集群上调度拓扑感知训练任务。此工作流程有助于提高训练实例放置的可见性和控制力。

在本文中，我们将逐步介绍查看节点拓扑信息并向集群提交拓扑感知任务。作为参考，NetworkNodes描述了实例的网络节点集。在每个网络节点集中，三个层构成了每个实例拓扑的分层视图。彼此最接近的实例将共享相同的第3层网络节点。如果在底层没有共同的网络节点，则查看第2层是否有共同性。

先决条件

要开始使用拓扑感知调度，您必须满足以下先决条件：

• EKS集群
• 启用了拓扑信息的SageMaker HyperPod集群
• 安装了SageMaker HyperPod任务治理附加组件（版本1.2.2或更高版本）
• 安装了Kubectl
• （可选）安装了SageMaker HyperPod CLI

获取节点拓扑信息

运行以下命令以显示集群中的节点标签。此命令提供每个实例的网络拓扑信息。

1
2
3

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

具有相同网络节点层3的实例尽可能接近，遵循EC2拓扑层次结构。您应该看到一个节点标签列表，类似于以下内容：topology.k8s.aws/network-node-layer-3: nn-33333example

运行以下脚本以显示集群中位于相同第1、2和3层网络节点上的节点：

1
2
3
4

git clone https://github.com/aws-samples/awsome-distributed-training.git
cd awsome-distributed-training/1.architectures/7.sagemaker-hyperpod-eks/task-governance 
chmod +x visualize_topology.sh
bash visualize_topology.sh

此脚本的输出将打印一个流程图，您可以在流图编辑器中使用它来可视化集群的节点拓扑。下图是一个七实例集群的集群拓扑示例。

提交任务

SageMaker HyperPod任务治理提供两种使用拓扑感知提交任务的方法。在本节中，我们讨论这两个选项和任务治理的第三个替代选项。

修改Kubernetes清单文件

首先，您可以修改现有的Kubernetes清单文件以包含以下两个注释选项之一：

• kueue.x-k8s.io/podset-required-topology - 如果必须将所有Pod调度在相同网络节点层上的节点上才能开始作业，请使用此选项
• kueue.x-k8s.io/podset-preferred-topology - 如果理想情况下希望将所有Pod调度在相同网络节点层上的节点上，但具有灵活性，请使用此选项

以下代码是使用kueue.x-k8s.io/podset-required-topology设置来调度共享相同第3层网络节点的Pod的示例作业：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

apiVersion: batch/v1
kind: Job
metadata:
  name: test-tas-job
  namespace: hyperpod-ns-team-a
  labels:
    kueue.x-k8s.io/queue-name: hyperpod-ns-team-a-localqueue
    kueue.x-k8s.io/priority-class: inference-priority
spec:
  parallelism: 10
  completions: 10
  suspend: true
  template:
    metadata:
      labels:
        kueue.x-k8s.io/queue-name: hyperpod-ns-team-a-localqueue
      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"]
          resources:
            requests:
              cpu: "1"
      restartPolicy: Never

要验证Pod正在哪些节点上运行，请使用以下命令查看每个Pod的节点ID：

1

kubectl get pods -n hyperpod-ns-team-a -o wide

使用SageMaker HyperPod CLI

第二种提交作业的方法是通过SageMaker HyperPod CLI。请确保安装最新版本以使用拓扑感知调度。要使用SageMaker HyperPod CLI进行拓扑感知调度，可以在创建作业命令中包含--preferred-topology参数或--required-topology参数。

以下代码是使用SageMaker HyperPod CLI启动拓扑感知mnist训练作业的示例命令，将XXXXXXXXXXXX替换为您的账户ID：

1
2
3
4
5
6
7

hyp create hyp-pytorch-job \
--job-name test-pytorch-job-cli \
--image XXXXXXXXXXXX.dkr.ecr.us-west-2.amazonaws.com/ptjob:mnist \
--pull-policy "Always" \
--tasks-per-node 1 \
--max-retry 1 \
--preferred-topology topology.k8s.aws/network-node-layer-3

清理

如果在遵循本文时部署了新资源，请参考SageMaker HyperPod EKS研讨会中的清理部分，以确保不会产生意外费用。

结论

在大语言模型训练期间，Pod到Pod通信将模型分布在多个实例上，需要这些实例之间频繁的数据交换。在本文中，我们讨论了SageMaker HyperPod任务治理如何通过优化吞吐量和延迟来帮助调度工作负载以实现作业效率。我们还逐步介绍了如何使用SageMaker HyperPod拓扑网络信息来调度作业，以优化AI任务的网络通信延迟。

我们鼓励您尝试此解决方案并在评论部分分享您的反馈。