通过 TAS 安排 GKE 工作负载

AI 和 ML 工作负载需要大量的 Pod 到 Pod 通信。由于此要求，Pod 之间的网络带宽会直接影响工作负载执行时间和费用。此带宽取决于集群中虚拟机 (VM) 实例的放置位置。

本文档介绍了如何在 Google Kubernetes Engine (GKE) 集群上优化大规模 AI 或机器学习工作负载的调度，以提高性能和可靠性。具体来说，您将配置集群以使用拓扑感知调度 (TAS) 实现低延迟通信。这种方法可最大限度地减少通信开销，并有助于最大限度地提高工作负载的性能。

什么是拓扑感知调度 (TAS)？

TAS 可以显著提高大语言模型 (LLM) 训练的效率。TAS 会在网络拓扑中战略性地放置工作器，以最大限度地减少梯度聚合期间的通信开销，这需要工作器按特定等级顺序进行通信。通过最大限度地减少按顺序通信的工作人员之间的网络跃点，TAS 可减少网络争用并优化带宽利用率，从而加快收敛速度并缩短训练时间。随着 LLM 模型越来越大，TAS 对于最大限度地提高分布式训练的性能和可伸缩性至关重要。

TAS 最适合使用密集放置的容量，而这种容量可以通过与预留绑定的容量来获得。使用 Spot 虚拟机和其他按用量付费选项时，您的容量不太可能紧密分配，因此 TAS 在这种情况下可能无法正常运行。

准备工作

在开始之前，请确保您已执行以下任务：

  • 启用 Google Kubernetes Engine API。
    • 启用 Google Kubernetes Engine API
  • 如果您要使用 Google Cloud CLI 执行此任务，请安装并初始化 gcloud CLI。如果您之前安装了 gcloud CLI，请通过运行 gcloud components update 命令来获取最新版本。较早版本的 gcloud CLI 可能不支持运行本文档中的命令。

  • 如需连接到集群，请运行以下命令：

    gcloud container clusters get-credentials CLUSTER_NAME

    将 CLUSTER_NAME 替换为您的集群名称。

准备 GKE 集群

如需准备 GKE 集群以运行使用 TAS 的工作负载，请完成以下步骤：

  1. 安装启用了 TAS 的 Kueue

  2. 查看 GKE 集群的拓扑

  3. 配置 Kueue

安装启用了 TAS 的 Kueue

我们建议您搭配使用 TAS 与 Kueue。Kueue 是一个 Kubernetes 原生系统，可用于管理配额以及作业使用配额的方式。TAS 需要使用 Kueue 版本 0.10.0 或更高版本，并且您必须明确启用它。

如需安装 Kueue 并启用 TAS，请选择以下选项之一：

Kueue 清单

  1. 安装 Kueue：

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/kueue/releases/download/v0.10.0/manifests.yaml

  2. 在 Kueue 中启用 TAS：

    kubectl -n kueue-system patch deployment kueue-controller-manager \
    --type json -p='[{"op": "add", "path": "/spec/template/spec/containers/0/args/-", "value": "--feature-gates=TopologyAwareScheduling=true"}]'

Helm 图表

使用 Helm 图表安装启用了 TAS 的 Kueue：

helm install kueue oci://us-central1-docker.pkg.dev/k8s-staging-images/charts/kueue \
    --version="v0.10.0" \
    --create-namespace \
    --namespace=kueue-system \
    --set="controllerManager.featureGates[0].name=TopologyAwareScheduling,controllerManager.featureGates[0].enabled=true"

安装 Kueue 后，您必须对其进行配置，使其了解其管理的相应基础设施，如下一部分中所述。

查看 GKE 集群的拓扑

在查看以抢占式虚拟机形式预配的 A4X、A4、A3 Ultra、A3 Mega 和 A3 High（8 个 GPU）节点的拓扑之前，您必须在 GKE 节点上定义紧凑放置，以便为 TAS 公开其物理拓扑。否则，您会遇到错误。

如需查看特定节点池中 GKE 集群节点的拓扑，请运行以下命令：

kubectl get nodes -l cloud.google.com/gke-nodepool=NODE_POOL_NAME \
    -ocustom-columns='NAME:.metadata.name,BLOCK:.metadata.labels.cloud\.google\.com/gce-topology-block,SUBBLOCK:.metadata.labels.cloud\.google\.com/gce-topology-subblock,HOST:.metadata.labels.cloud\.google\.com/gce-topology-host' | sort -k2,4

将 NODE_POOL_NAME 替换为节点池的名称。

如需了解输出中虚拟机上 GKE 节点的物理拓扑，请参阅以下节点标签：

  • cloud.google.com/gce-topology-block：虚拟机所在预留块的组织专用 ID。

  • cloud.google.com/gce-topology-subblock：虚拟机所在子块的组织专用 ID。