AI 和機器學習工作負載需要大量的 Pod 對 Pod 通訊。因此，Pod 之間的網路頻寬會直接影響工作負載執行時間和成本。這項頻寬取決於叢集中虛擬機器 (VM) 執行個體的位置。
本文說明如何最佳化 Google Kubernetes Engine (GKE) 叢集的大規模 AI 或機器學習工作負載排程，兼顧效能和可靠性。具體來說，您會將叢集設為使用拓撲感知調度 (TAS)，以進行低延遲通訊。這種做法可盡量減少通訊負擔，並盡可能提高工作負載效能。
什麼是拓撲感知調度 (TAS)？
TAS 可大幅提升大型語言模型 (LLM) 的訓練效率。TAS 會根據 Network Topology 策略性地放置 worker，盡可能降低梯度匯總期間的通訊負擔，這項作業需要 worker 按照特定排名順序通訊。TAS 可盡量減少依序通訊 worker 之間的網路躍點，進而減少網路爭用情形，並充分運用頻寬，加快收斂速度並縮短訓練時間。隨著 LLM 模型越來越大，TAS 對於最大化分散式訓練的效能和擴充性至關重要。
TAS 最適合搭配密集放置的容量，這類容量可透過預留繫結容量取得。如果使用 Spot VM 和其他用量選項，容量較不容易同時配置，因此 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 執行工作負載，請完成下列步驟：
安裝 Kueue 並啟用 TAS
建議您搭配使用 TAS 與 Kueue，後者是 Kubernetes 原生系統，可管理配額和工作耗用配額的方式。TAS 必須使用 Kueue 0.10.0 以上版本，且須明確啟用。
如要安裝 Kueue 並啟用 TAS，請選取下列其中一個選項：
Kueue 資訊清單
安裝 Kueue：
kubectl apply --server-side -f https://github.com/kubernetes-sigs/kueue/releases/download/v0.10.0/manifests.yaml在 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 chart 安裝 Kueue，並啟用 TAS：
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 後，您必須設定 Kueue，讓它瞭解所管理的基礎架構，詳情請參閱下一節。
查看 GKE 叢集的拓撲
如要查看 A4X、A4、A3 Ultra、A3 Mega 和 A3 High (8 個 GPU) 節點 (佈建為 Spot VM) 的拓撲，您必須先在 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 替換為節點集區名稱。
如要瞭解輸出內容中 VM 上 GKE 節點的實體拓撲，請參閱下列節點標籤：
cloud.google.com/gce-topology-block：VM 所在預留區塊的組織專屬 ID。cloud.google.com/gce-topology-subblock：VM 所在子區塊的組織專屬 ID。cloud.google.com/gce-topology-host：VM 所在主機的 ID。kubernetes.io/hostname：Kubernetes 節點的主機名稱。這個主機名稱通常也是 GKE 節點名稱。
兩個 VM 共用的標籤值越多，實體位置就越靠近。如要進一步瞭解這些詞彙，請參閱「術語」。
設定 Kueue
安裝 Kueue 後，您必須設定 Kueue，指定要管理的基礎架構。一般來說，Kueue 需要靜態基礎架構或啟用叢集自動調度資源功能的動態基礎架構，才能定義 ClusterQueue 資源配額。只有在工作負載要求的資源小於或等於 ClusterQueue 中定義的資源集區時，ClusterQueue 才會接受工作負載。按照本節說明設定 Kueue 後，Kueue 會使用 TAS 准予工作負載，如下所示：
TAS 工作負載：Kueue 會檢查實體基礎架構的拓撲和目前用量。
非 TAS 工作負載：Kueue 不會檢查實體基礎架構的拓撲。Kueue 會管理設定中定義的所有配額，並將節點指派作業交給 kube-scheduler。
如要瞭解如何向 Kueue 提供 ClusterQueue 資源配額定義，請參閱下列範例：
配額極高：Kueue 幾乎不會根據要求的資源停止准許工作負載。根據 TAS 定義，Kueue 可能會依據基礎架構拓撲，決定是否允許工作負載。詳情請參閱「極高資源配額」。
實際配額：只有在工作負載要求的資源不超出這些資源配額限制時，Kueue 才會允許工作負載。Kueue 會根據 TAS 定義，先檢查基礎架構拓撲，再允許工作負載。詳情請參閱「實際資源配額」。
下列各節中提及的資源配額，均指 ClusterQueue 資源配額。
資源配額極高
下列範例使用非常高的資源配額，因此 Kueue 絕不會根據可用資源配額停止工作負載。而是使用可用節點的拓撲資訊，嘗試比對拓撲與工作負載的需求。
如要使用下列資源配額定義，請完成下列步驟：
開啟您偏好的檔案編輯器。接著，在名為
kueue-tas-config-very-high-quota.yaml的 YAML 檔案中加入下列配額定義：apiVersion: kueue.x-k8s.io/v1alpha1 kind: Topology metadata: name: "gke-default" spec: levels: - nodeLabel: "cloud.google.com/gce-topology-block" - nodeLabel: "cloud.google.com/gce-topology-subblock" - nodeLabel: "cloud.google.com/gce-topology-host" - nodeLabel: "kubernetes.io/hostname" --- kind: ResourceFlavor apiVersion: kueue.x-k8s.io/v1beta1 metadata: name: "tas-flavor" spec: nodeLabels: cloud.google.com/gke-nodepool: "NODE_POOL_NAME" topologyName: "gke-default" tolerations: - key: "nvidia.com/gpu" operator: "Exists" effect: NoSchedule --- apiVersion: kueue.x-k8s.io/v1beta1 kind: ClusterQueue metadata: name: "tas-cluster-queue" spec: namespaceSelector: {} resourceGroups: - coveredResources: ["nvidia.com/gpu"] flavors: - name: "tas-flavor" resources: - name: "nvidia.com/gpu" nominalQuota: 10000000 --- apiVersion: kueue.x-k8s.io/v1beta1 kind: LocalQueue metadata: namespace: "default" name: "tas-user-queue" spec: clusterQueue: "tas-cluster-queue"將
NODE_POOL_NAME替換為節點集區名稱。建立並套用 Kueue 工作佇列系統的資源配額設定：
kubectl create -f kueue-tas-config-very-high-quota.yaml
實際資源配額
上一個範例只設定了 GPU 資源。不過，Kueue 可以管理所有與 Kubernetes 相容的資源。
以下範例定義了更符合實際情況的資源配額，包括 CPU、記憶體和 GPU。這是 100 部 a3-ultragpu-8g 機器的費用。單一機器有 224 個 vCPU、2944 GB 記憶體和 8 個 GPU。同樣的原則也適用於一般 GPU (A2、A3 Edge、G2、G4 和 N1) 機型，使用各自的 CPU 和記憶體規格。
如要使用下列資源配額定義，請完成下列步驟：
開啟您偏好的檔案編輯器。接著，在名為
kueue-tas-config-real-quota.yaml的 YAML 檔案中加入下列配額定義：apiVersion: kueue.x-k8s.io/v1alpha1 kind: Topology metadata: name: "gke-default" spec: levels: - nodeLabel: "cloud.google.com/gce-topology-block" - nodeLabel: "cloud.google.com/gce-topology-subblock" - nodeLabel: "cloud.google.com/gce-topology-host" - nodeLabel: "kubernetes.io/hostname" --- kind: ResourceFlavor apiVersion: kueue.x-k8s.io/v1beta1 metadata: name: "tas-flavor" spec: nodeLabels: cloud.google.com/gke-nodepool: "NODE_POOL_NAME" topologyName: "gke-default" tolerations: - key: "nvidia.com/gpu" operator: "Exists" effect: NoSchedule --- apiVersion: kueue.x-k8s.io/v1beta1 kind: ClusterQueue metadata: name: "tas-cluster-queue" spec: namespaceSelector: {} # match all resourceGroups: - coveredResources: ["cpu", "memory", "nvidia.com/gpu"] flavors: - name: "tas-flavor" resources: # numbers below represent quota of 100 a3-ultragpu-8g machines - name: "cpu" nominalQuota: 22400 - name: "memory" nominalQuota: 294400Gi - name: "nvidia.com/gpu" nominalQuota: 800 --- apiVersion: kueue.x-k8s.io/v1beta1 kind: LocalQueue metadata: namespace: "default" name: "tas-user-queue" spec: clusterQueue: "tas-cluster-queue"將
NODE_POOL_NAME替換為節點集區名稱。建立資源配額設定，並套用至 Kueue 工作佇列系統：
kubectl create -f kueue-tas-config-real-quota.yaml輸出結果會與下列內容相似：
topology.kueue.x-k8s.io/gke-default created resourceflavor.kueue.x-k8s.io/tas-flavor created clusterqueue.kueue.x-k8s.io/tas-cluster-queue created localqueue.kueue.x-k8s.io/tas-user-queue created
使用 Kueue 透過 TAS 排定工作負載
下列情境說明如何指示 Kueue 和 TAS，使用拓撲要求類型和拓撲要求層級，管理常見的工作負載和基礎架構組合：
可用的拓撲要求類型如下 (偏好或必要)：
kueue.x-k8s.io/podset-preferred-topology：Kueue 會優先在指定拓撲層級內排定整個工作負載，但仍會接受不符合這個拓撲層級的工作負載。對於可能適合單一拓撲層級的工作負載，Kueue 可能會跨該拓撲層級的多個執行個體排程該工作負載。kueue.x-k8s.io/podset-required-topology：Kueue 會持續嘗試允許這項工作負載，直到整個工作負載都能符合所選拓撲層級。
以下是可用的拓撲要求層級，可讓您更具體或更籠統地指定偏好或要求執行工作的實體基礎架構：
cloud.google.com/gce-topology-blockcloud.google.com/gce-topology-subblockcloud.google.com/gce-topology-hostkubernetes.io/hostname
如要使用這些值排定工作負載，請使用下列 Job YAML 檔案：
apiVersion: batch/v1
kind: Job
metadata:
generateName: JOB_NAME
labels:
kueue.x-k8s.io/queue-name: tas-user-queue
spec:
parallelism: NUMBER_OF_REPLICAS
completions: NUMBER_OF_REPLICAS
completionMode: Indexed
template:
metadata:
annotations:
ANNOTATIONS_STRING
spec:
containers:
- name: dummy-job
image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
args: ["60s"]
resources:
requests:
nvidia.com/gpu: "1"
limits:
nvidia.com/gpu: "1"
restartPolicy: Never
請替換下列變數：
JOB_NAME：工作名稱。NUMBER_OF_REPLICAS：平行執行的 Pod 數量。ANNOTATIONS_STRING：請參閱下表：要求的拓撲類型和層級 說明 ANNOTATIONS_STRING偏好在主機名稱內執行 (建議) 只要有足夠資源可滿足工作負載的資源需求，即使容量分散，這項設定也會允許工作負載。Kueue 會盡可能緊密地排定 Pod。 kueue.x-k8s.io/podset-preferred-topology: "kubernetes.io/hostname"必須在主機中執行 只有在主機有足夠資源滿足工作負載的資源需求時，這項設定才會允許工作負載。
如果每部主機有多個 VM (例如較小的機器類型)，或多個 Pod 可以在單一節點上執行，這項功能就非常實用。在這種情況下，如果工作負載獲准，就會在單一主機上執行。
kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-host"偏好在主機中執行 只要有足夠資源可滿足工作負載的資源需求，即使容量分散，這項設定也會允許工作負載。Kueue 會嘗試在主機中排定 Pod，並視需要使用其他主機。 kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-host"必須在子模塊中執行 只有在子區塊有足夠資源滿足工作負載的資源需求時，這項設定才會允許工作負載。 kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-subblock"建議在子區塊內執行 只要有足夠資源可滿足工作負載的資源需求，即使容量分散，這項設定也會允許工作負載。Kueue 會嘗試在子區塊中排定 Pod，並視需要使用其他子區塊。在這種情況下，Kueue 會優先選擇可用容量較多的子區塊，即使該子區塊已分割，也會優先於容量剛好滿足需求的子區塊。 kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-subblock"必須在區塊內執行 這項設定會檢查區塊內的可用資源是否符合工作負載的資源需求，只有在符合時才會允許工作負載。如果准入，Kueue 會盡可能減少子區塊和主機數量，以便排定工作負載。這可能會導致可用容量分散。 kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"建議在區塊內執行 只要有足夠資源可滿足工作負載的資源需求，即使容量分散，這項設定也會允許工作負載。Kueue 會嘗試在一個區塊內排定 Pod，並視需要使用其他區塊。 kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-block"
使用 Kueue 搭配 TAS，透過 PodGroup 排定工作負載
使用 PodGroups 時，您必須為 PodGroup 中的每個 Pod 指定三個額外欄位：
標籤：
kueue.x-k8s.io/pod-group-name：用於匯總的 PodGroup 名稱。
kueue.x-k8s.io/pod-group-pod-index：PodGroup 中每個 Pod 的索引。
註解：
- kueue.x-k8s.io/pod-group-total-count：PodGroup 中的 Pod 總數。
視您使用的機器學習架構而定，PodGroup 的領導者可能需要 GPU，也可能不需要。由於 Kueue 的限制，這些情況需要以不同方式處理。下列範例說明如何建立 PodGroup，其中包含三個 Pod (一個 leader 和兩個 worker)。
案例 1：Leader 也是 worker，且需要 GPU
如果 leader 是其中一個 worker，且需要 GPU，則 leader 在 PodGroup 中可以有任何數量。為求簡單，在以下範例中，領袖的索引為 0：
apiVersion: v1
kind: Pod
metadata:
generateName: tas-podgroup-leader-
labels:
kueue.x-k8s.io/queue-name: tas-user-queue
kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
kueue.x-k8s.io/pod-group-pod-index: "0"
annotations:
kueue.x-k8s.io/pod-group-total-count: "3"
kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
containers:
- name: leader
image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
args: ["600s"]
resources:
requests:
nvidia.com/gpu: "1"
limits:
nvidia.com/gpu: "1"
restartPolicy: Never
---
apiVersion: v1
kind: Pod
metadata:
generateName: tas-podgroup-worker-1-
labels:
kueue.x-k8s.io/queue-name: tas-user-queue
kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
kueue.x-k8s.io/pod-group-pod-index: "1"
annotations:
kueue.x-k8s.io/pod-group-total-count: "3"
kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
restartPolicy: Never
containers:
- name: worker
image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
args: ["600s"]
resources:
requests:
nvidia.com/gpu: "1"
limits:
nvidia.com/gpu: "1"
---
apiVersion: v1
kind: Pod
metadata:
generateName: tas-podgroup-worker-2-
labels:
kueue.x-k8s.io/queue-name: tas-user-queue
kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
kueue.x-k8s.io/pod-group-pod-index: "2"
annotations:
kueue.x-k8s.io/pod-group-total-count: "3"
kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
restartPolicy: Never
containers:
- name: worker
image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
args: ["600s"]
resources:
requests:
nvidia.com/gpu: "1"
limits:
nvidia.com/gpu: "1"
案例 2：Leader 不是 worker，不需要 GPU
如果因 Kueue 限制而導致 leader 不是 worker 之一，則 leader 必須在 PodGroup 中具有最後一個索引，因為 Kueue 建立 PodSet 的方式。如果 leader 沒有最後一個索引，且第一個 worker 未使用第一個索引，Kueue 就不會套用排名指派。
請參閱以下範例：
---
apiVersion: v1
kind: Pod
metadata:
generateName: tas-podgroup-leader-
labels:
kueue.x-k8s.io/queue-name: tas-user-queue
kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
kueue.x-k8s.io/pod-group-pod-index: "2"
annotations:
kueue.x-k8s.io/pod-group-total-count: "3"
kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
containers:
- name: leader
image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
args: ["600s"]
resources:
requests:
cpu: "1"
limits:
cpu: "1"
restartPolicy: Never
---
apiVersion: v1
kind: Pod
metadata:
generateName: tas-podgroup-worker-0-
labels:
kueue.x-k8s.io/queue-name: tas-user-queue
kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
kueue.x-k8s.io/pod-group-pod-index: "0"
annotations:
kueue.x-k8s.io/pod-group-total-count: "3"
kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
restartPolicy: Never
containers:
- name: worker
image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
args: ["600s"]
resources:
requests:
nvidia.com/gpu: "1"
limits:
nvidia.com/gpu: "1"
---
apiVersion: v1
kind: Pod
metadata:
generateName: tas-podgroup-worker-1-
labels:
kueue.x-k8s.io/queue-name: tas-user-queue
kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
kueue.x-k8s.io/pod-group-pod-index: "1"
annotations:
kueue.x-k8s.io/pod-group-total-count: "3"
kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
restartPolicy: Never
containers:
- name: worker
image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
args: ["600s"]
resources:
requests:
nvidia.com/gpu: "1"
limits:
nvidia.com/gpu: "1"
後續步驟
如要進一步瞭解如何在 GKE 叢集中啟用節點健康狀態預測功能，請參閱「啟用節點健康狀態預測功能」。
如要瞭解如何管理與 GKE 叢集和 AI 工作負載相關的常見事件，請參閱「管理 AI 最佳化 GKE 叢集」。
如要進一步瞭解如何使用 Kueue 在 GKE 上排定工作，請參閱「使用 Kueue 部署批次系統」。
如要使用預先設定的 Kueue 和 TAS 藍圖部署 GKE 叢集，請參閱「Cluster Toolkit 中的 GKE 部署作業總覽」。