创建使用 A4X 的自定义 AI 优化型 GKE 集群

本页面介绍了如何创建经过 AI 优化的 Google Kubernetes Engine (GKE) 集群，该集群使用 A4X 虚拟机 (VM) 来支持 AI 和机器学习工作负载。如需详细了解 A4X，请参阅 A4X 系列。

A4X 机器系列旨在让您能够运行大规模 AI/机器学习集群，并提供有针对性的工作负载放置、高级集群维护控制和拓扑感知调度等功能。如需了解详情，请参阅集群管理概览。

GKE 提供了一个统一的平台界面，可用于运行组织中的各种工作负载，从而减轻管理多个平台的运营负担。您可以运行高性能分布式预训练、模型微调、模型推理、应用服务和支持服务等工作负载。

本页面介绍如何使用 Google Cloud CLI 创建 GKE 集群，以便根据工作负载的需求尽可能灵活地配置集群。或者，您也可以选择使用 Cluster Toolkit 快速部署集群，其默认设置反映了许多使用情形下的最佳实践。

如需了解详情，请参阅创建具有默认配置的 AI 优化型 GKE 集群。
如需创建使用 A4 或 A3 Ultra 的集群，请参阅创建使用 A4 或 A3 Ultra 的自定义 AI 优化型 GKE 集群。

启用 GPUDirect RDMA 时的集群配置选项

如需使用 Google Cloud CLI 创建集群，您可以选择以下任一集群配置选项：

如果您计划运行分布式 AI 工作负载：请按照本页中的说明创建具有 GPUDirect RDMA 的 GKE 集群。
如果您不打算运行分布式 AI 工作负载：请创建一个不使用 GPUDirect RDMA 的 GKE 集群。如需了解详情，请参阅创建不含 GPUDirect RDMA 的集群。

准备工作

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

启用 Google Kubernetes Engine API。

启用 Google Kubernetes Engine API

如果您要使用 Google Cloud CLI 执行此任务，请安装并初始化 gcloud CLI。如果您之前安装了 gcloud CLI，请通过运行 gcloud components update 命令来获取最新版本。较早版本的 gcloud CLI 可能不支持运行本文档中的命令。
注意：对于现有 gcloud CLI 安装，请务必设置 compute/region 属性。如果您主要使用可用区级集群，请改为设置 compute/zone。通过设置默认位置，您可以避免 gcloud CLI 中出现如下错误：One of [--zone, --region] must be supplied: Please specify location。如果集群的位置与您设置的默认位置不同，您可能需要在某些命令中指定位置。

所需的角色

如需获得创建和管理 GKE 集群所需的权限，请让您的管理员为您授予项目的以下 IAM 角色：

Kubernetes Engine Admin (roles/container.admin)
Cloud Build Editor (roles/cloudbuild.builds.editor)
Compute Admin (roles/compute.admin)
Project IAM Admin (roles/resourcemanager.projectIamAdmin)
Service Account Admin (roles/iam.serviceAccountAdmin)
Service Account User (roles/iam.serviceAccountUser)
Service Usage Consumer (roles/serviceusage.serviceUsageConsumer)
Storage Admin (roles/storage.admin)

如需详细了解如何授予角色，请参阅管理对项目、文件夹和组织的访问权限。

您也可以通过自定义角色或其他预定义角色来获取所需的权限。

获取容量

您可以通过创建未来预留来获取 A4X 虚拟机的容量。如需详细了解未来预留，请参阅选择使用选项表格中的AI Hypercomputer 中的未来预留列。

如需通过未来预留获取容量，请参阅如何获取容量表格中的“AI Hypercomputer 中的未来预留”行。

要求

AI 优化型 GKE 集群（使用 A4X 虚拟机）需满足以下要求：

对于 A4X，请确保您使用的是 GKE 1.33.4-gke.1036000 版或更高版本（适用于 1.33 或更高版本）。或者，对于 1.32，请使用 GKE 1.32.8-gke.1108000 版或更高版本。这些版本可确保 A4X 使用以下内容：
- R580，A4X 的最低 GPU 驱动程序版本。
- 基于相干驱动程序的内存管理 (CDMM)，默认情况下处于启用状态。NVIDIA 建议 Kubernetes 集群启用此模式，以解决内存报告过多的问题。借助 CDMM，可以通过驱动程序（而非操作系统 [OS]）来管理 GPU 内存。此方法有助于避免操作系统在线使用 GPU 内存，并将 GPU 内存作为非统一内存访问 (NUMA) 节点公开给操作系统。启用 CDMM 时，不支持多实例 GPU。如需详细了解 CDMM，请参阅硬件和软件支持。
- GPUDirect RDMA，建议启用此功能，以便 A4X 节点池使用 A4X 的网络功能。
GKE 节点必须使用 Container-Optimized OS 节点映像。不支持 Ubuntu 和 Windows 节点映像。
您的 GKE 工作负载必须使用所有可用的 GPU，并且您的 Pod 必须使用单个 GKE 节点上所有可用的次要 NIC。多个 Pod 无法在单个 GKE 节点上共享 RDMA。
此设置会运行 NCCL 测试。如需运行此 NCCL 测试，您必须至少有 2 个虚拟机的配额（使用 a4x-highgpu-4g 或 a4x-highgpu-4g-nolssd 时，每个虚拟机有 4 个 GPU）。
您必须使用受预留约束的预配模型来创建包含 A4X 的集群。不支持其他配置模型。
GPUDirect RDMA 与 NCCL Fast Socket 或 GPUDirect TCPX/TCPXO 不兼容。请勿在使用 GPUDirect RDMA 的集群上启用 NCCL Fast Socket 或安装 TCPX/TCPXO 插件。

创建集群的注意事项

创建集群时，请考虑以下信息：

选择集群位置：
- 验证您使用的位置是否提供您选择的机器类型。如需了解详情，请参阅 GPU 位置。
- 对于密集型预留，您可以创建区域级集群。在这种情况下，请将 --region 标志替换为 --zone=COMPUTE_ZONE 标志，其中 COMPUTE_ZONE 是控制平面的可用区。
- 在区域级集群中创建节点池时，您可以使用 --node-locations 标志为 GKE 节点指定可用区。
选择驱动程序版本：
- 驱动程序版本可以是以下值之一：
  - default：为您的 GKE 节点版本安装默认驱动程序版本。如需详细了解默认驱动程序版本的相关要求，请参阅要求部分。
  - latest：为您的 GKE 版本安装最新可用的驱动程序版本。此选项仅适用于使用 Container-Optimized OS 的节点。
  - disabled：跳过自动驱动程序安装。创建节点池后，您必须手动安装驱动程序。
- 如需详细了解 GKE 节点版本的默认和最新 GPU 驱动程序版本，请参阅手动安装 NVIDIA GPU 驱动程序部分中的表格。
选择预留亲和性：
- 您可以找到有关预留的信息，例如预留的名称或预留中特定块的名称。如需查找这些值，请参阅查看未来预留请求。
- --reservation-affinity 标志的值可以是 specific 或 any。不过，对于高性能分布式 AI 工作负载，我们建议您使用特定预留。
- 使用特定预留（包括共享预留）时，请按以下格式指定 --reservation 标志的值：
```
projects/PROJECT_ID/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME
```
  替换以下值：
  - PROJECT_ID：您的 Google Cloud 项目 ID。
  - RESERVATION_NAME：预留的名称。
  - BLOCK_NAME：预留中特定块的名称。
  如需使用以子块为目标的预留，以便将虚拟机放置在 BLOCK_NAME 内的单个子块中，请将以下内容添加到路径末尾：
```
/reservationSubBlocks/SUB_BLOCK_NAME
```
  将 SUB_BLOCK_NAME 替换为子块的名称。
提示：如果预留位于当前项目中，则可以从预留值中省略 projects/PROJECT_ID/reservations/。

创建使用 A4X 和 GPUDirect RDMA 的 AI 优化型 GKE 集群

对于分布式 AI 工作负载，多个 GPU 节点通常会链接在一起，作为一个计算机运行。A4X 是基于 NVIDIA GB200 NVL72 机架级架构的艾级平台。此机器类型可为 AI 工作负载提供高性能的云体验，从而实现跨多个 GPU 的伸缩和协作。如需详细了解 A4X 的网络架构（包括网络带宽和 NIC 配置），请参阅 A4X 机器类型。

如需创建支持 A4X 和 GPUDirect RDMA 的 GKE Standard 集群，请完成以下步骤（将在后续部分中介绍）：

创建 VPC 和子网
创建具有多网络功能的 GKE 集群
创建 GKE 网络对象
创建工作负载政策
创建具有 A4X 的节点池
安装 RDMA 二进制文件并配置 NCCL
安装 NVIDIA Compute Domain CRD 和 DRA 驱动程序

创建 VPC 和子网

A4X 虚拟机具有以下配置：

每个虚拟机配备 4 个 NVIDIA B200 GPU，通过 NVLink 连接
两个基于 Arm 的 NVIDIA Grace CPU
四个 400 Gbps CX-7 网络接口卡 (NIC)，用于 GPU 到 GPU 的联网
两个 200 Gbps Google Titanium 网络接口卡 (NIC)，用于外部服务

AI 和 ML 工作负载（例如分布式训练）需要通过缩短作业完成时间来获得极大的加速，从而优化性能。对于需要高性能、高吞吐量和低延迟的工作负载，GPUDirect RDMA 可减少将载荷传输到 GPU 和从 GPU 传输载荷所需的网络跃点数。这种方法可以更高效地利用可用的网络带宽。

与 CPU 关联的某个 Google Titanium NIC 使用 GKE 中的默认网络，因此只要默认网络有足够的 IP 地址范围，您就不必为此 NIC 创建新的 VPC。

您可以使用以下命令为第二个 CPU Titanium NIC (gVNIC) 创建一个 VPC，并为四个 CX-7 RDMA NIC 创建另一个 VPC。

为了最大限度地提高网络带宽，用于为额外的 GVNIC 创建 VPC 的命令将最大传输单元 (MTU) 设置为 8896。RDMA VPC 默认采用建议的 8896 设置。如需了解详情，请参阅 MTU 设置和 GPU 机器类型。

设置环境变量以匹配您的部署：
```
export REGION="COMPUTE_REGION"
export ZONE="COMPUTE_ZONE"
export PROJECT="PROJECT_ID"
export GVNIC_NETWORK_PREFIX="GVNIC_NETWORK_PREFIX"
export RDMA_NETWORK_PREFIX="RDMA_NETWORK_PREFIX"
```
执行以下变量替换操作：
- COMPUTE_REGION：集群所在的区域。
- COMPUTE_ZONE：节点池的可用区。
- PROJECT_ID：您的 Google Cloud 项目 ID。
- GVNIC_NETWORK_PREFIX：GVNIC 网络前缀（例如 a4x-gvnic）。
- RDMA_NETWORK_PREFIX：RDMA 网络前缀（例如 a4x-rdma）。

创建两个 VPC 网络：

# Create a VPC for the additional GVNIC
gcloud compute --project=${PROJECT} \
  networks create \
  GVNIC_NETWORK_PREFIX-net \
  --subnet-mode=custom \
  --mtu=8896

gcloud compute --project=${PROJECT} \
  networks subnets create \
  GVNIC_NETWORK_PREFIX-sub \
  --network=GVNIC_NETWORK_PREFIX-net \
  --region=${REGION} \
  --range=192.168.0.0/24

gcloud compute --project=${PROJECT} \
  firewall-rules create \
  GVNIC_NETWORK_PREFIX-internal \
  --network=GVNIC_NETWORK_PREFIX-net \
  --action=ALLOW \
  --rules=tcp:0-65535,udp:0-65535,icmp \
  --source-ranges=192.168.0.0/16

# Create HPC VPC for the RDMA NICs with 4 subnets.
gcloud compute --project=${PROJECT} \
  networks create RDMA_NETWORK_PREFIX-net \
  --network-profile=${ZONE}-vpc-roce \
  --subnet-mode=custom

# Create subnets for the HPC VPC.
for N in $(seq 0 3); do
  gcloud compute --project=${PROJECT} \
    networks subnets create \
    RDMA_NETWORK_PREFIX-sub-$N \
    --network=RDMA_NETWORK_PREFIX-net \
    --region=${REGION} \
    --range=192.168.$((N+1)).0/24 &  # offset to avoid overlap with gvnics
done

创建具有多网络功能的 GKE 集群

创建具有多网络功能的 GKE Standard 集群：
```
gcloud container clusters create CLUSTER_NAME \
    --enable-dataplane-v2 --enable-ip-alias --location=COMPUTE_REGION \
    --enable-multi-networking --cluster-version=CLUSTER_VERSION \
    --enable-kubernetes-unstable-apis=resource.k8s.io/v1beta1/deviceclasses,resource.k8s.io/v1beta1/resourceclaims,resource.k8s.io/v1beta1/resourceclaimtemplates,resource.k8s.io/v1beta1/resourceslices \
    [--services-ipv4-cidr=SERVICE_CIDR \
    --cluster-ipv4-cidr=POD_CIDR]
```
准备此命令时，请执行以下操作：
- 创建运行 GKE 1.34 或更高版本的集群时，请移除用于启用 DRA 的 --enable-kubernetes-unstable-apis 标志。只有在创建运行 1.32 版或 1.33 版的集群时，才需要使用此标志。启用 Beta 版 API 后，当该 API 在后续 Kubernetes 次要版本中被弃用并移除时，您最终必须从该 API 迁移。如需了解详情，请参阅 Beta 版 API 弃用。
- 替换以下内容：
  - CLUSTER_NAME：您的集群的名称。
  - CLUSTER_VERSION：新集群的版本。如需详细了解哪个版本的 GKE 支持您的配置，请参阅本文档中的要求。
  - COMPUTE_REGION：计算区域的名称。
- （可选）您可以明确指定服务和 Pod 的次要 CIDR 范围。如果您使用这些可选标志，请替换以下变量：
  - SERVICE_CIDR：服务的次要 CIDR 范围。
  - POD_CIDR：Pod 的次要 CIDR 范围。
  使用这些标志时，您必须验证 CIDR 范围是否与额外节点网络的子网范围重叠。例如，考虑 SERVICE_CIDR=10.65.0.0/19 和 POD_CIDR=10.64.0.0/19。如需了解详情，请参阅添加 Pod IPv4 地址范围。

创建 GKE 网络对象

您必须通过 GKE 网络参数集配置上一部分中创建的 VPC 网络。具体而言，第二个 CPU Titanium NIC (gVNIC) 需要配置为 NetDevice 模式，而四个 CX-7 RDMA NIC 都需要配置为 RDMA 模式。

此命令使用以下名称：

CPU Titanium NIC (gVNIC) VPC 的名称为 GVNIC_NETWORK_PREFIX-net，子网的名称为 GVNIC_NETWORK_PREFIX-sub
CX-7 RDMA NIC 的 VPC 名为 RDMA_NETWORK_PREFIX-net，子网名为 RDMA_NETWORK_PREFIX-sub-[0…3]

运行以下命令以创建 GKE 网络对象：

kubectl apply -f - <<EOF
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: gvnic-1
spec:
  vpc: GVNIC_NETWORK_PREFIX-net
  vpcSubnet: GVNIC_NETWORK_PREFIX-sub
  deviceMode: NetDevice
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: gvnic-1
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: gvnic-1
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-0
spec:
  vpc: RDMA_NETWORK_PREFIX-net
  vpcSubnet: RDMA_NETWORK_PREFIX-sub-0
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-0
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-0
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-1
spec:
  vpc: RDMA_NETWORK_PREFIX-net
  vpcSubnet: RDMA_NETWORK_PREFIX-sub-1
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-1
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-1
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-2
spec:
  vpc: RDMA_NETWORK_PREFIX-net
  vpcSubnet: RDMA_NETWORK_PREFIX-sub-2
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-2
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-2
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: rdma-3
spec:
  vpc: RDMA_NETWORK_PREFIX-net
  vpcSubnet: RDMA_NETWORK_PREFIX-sub-3
  deviceMode: RDMA
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: rdma-3
spec:
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: rdma-3
EOF

创建工作负载政策

必须有工作负载政策才能创建分区。如需了解详情，请参阅 MIG 的工作负载政策。

创建 HIGH_THROUGHPUT 工作负载政策，并将 accelerator_topology 字段设置为 1x72。

gcloud beta compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
    --type HIGH_THROUGHPUT \
    --accelerator-topology 1x72 \
    --project PROJECT \
    --region COMPUTE_REGION

替换以下内容：

WORKLOAD_POLICY_NAME：工作负载政策的名称。
PROJECT：您的项目的名称。
COMPUTE_REGION：计算区域的名称。

创建具有 A4X 的节点池

我们建议您创建一个使用 GKE GPU 设备插件的节点池。此插件提供 GKE 管理的 GPU 资源管理。此方法具有以下优势：

易于部署和升级
驱动程序自动安装
GKE 管理的 GPU 功能，例如指标和分区 GPU
基本安全漏洞修复

或者，如果您的使用情形需要，您也可以使用 NVIDIA GPU Operator。如需了解详情，请参阅为什么使用 NVIDIA GPU Operator？。

使用 GKE GPU 设备插件创建 A4X 节点池

创建使用 GKE GPU 设备插件的 A4X 节点池：

gcloud container node-pools create NODE_POOL_NAME \
    --location COMPUTE_REGION \
    --node-locations COMPUTE_ZONE \
    --cluster CLUSTER_NAME \
    --num-nodes=NODE_COUNT \
    --machine-type MACHINE_TYPE \
    --accelerator type=nvidia-gb200,count=4,gpu-driver-version=DRIVER_VERSION \
    --additional-node-network network=GVNIC_NETWORK_PREFIX-net,subnetwork=GVNIC_NETWORK_PREFIX-sub \
    --additional-node-network network=RDMA_NETWORK_PREFIX-net,subnetwork=RDMA_NETWORK_PREFIX-sub-0 \
    --additional-node-network network=RDMA_NETWORK_PREFIX-net,subnetwork=RDMA_NETWORK_PREFIX-sub-1 \
    --additional-node-network network=RDMA_NETWORK_PREFIX-net,subnetwork=RDMA_NETWORK_PREFIX-sub-2 \
    --additional-node-network network=RDMA_NETWORK_PREFIX-net,subnetwork=RDMA_NETWORK_PREFIX-sub-3 \
    --scopes "https://www.googleapis.com/auth/cloud-platform" \
    --reservation-affinity=specific \
    --reservation=RESERVATION_NAME/reservationBlocks/BLOCK_NAME \
    --placement-policy=WORKLOAD_POLICY_NAME

替换以下内容：

NODE_POOL_NAME：节点池的名称。
COMPUTE_REGION：集群所在的区域。
COMPUTE_ZONE：节点池的可用区。
CLUSTER_NAME：您的集群的名称。
NODE_COUNT：节点池的节点数，必须为 18 个或更少。我们建议使用 18 个节点，以便通过 NVLink 网域在一个子块中获得 1x72 的 GPU 拓扑。
MACHINE_TYPE：a4x-highgpu-4g 或 a4x-highgpu-4g-nolssd，具体取决于您是否需要本地 SSD。
DRIVER_VERSION：要安装的 NVIDIA 驱动程序版本。它可以是以下值之一：default、latest 或 disabled。
RESERVATION_NAME：预留的名称。如需查找此值，请参阅查看未来预留请求。
BLOCK_NAME：预留中特定块的名称。如需查找此值，请参阅查看未来预留请求。
WORKLOAD_POLICY_NAME：您之前创建的工作负载政策的名称。

使用 NVIDIA GPU Operator 创建 A4X 节点池

或者，如需使用 NVIDIA GPU Operator，请执行以下步骤：

运行上一部分中的 gcloud container node-pools create 命令，并进行以下更改：
- 将 gpu-driver-version=latest 更改为 gpu-driver-version=disabled。此修改会跳过自动 GPU 驱动程序安装，因为使用 NVIDIA GPU Operator 时不支持自动安装。
- 设置 --node-labels="gke-no-default-nvidia-gpu-device-plugin=true" 以停用 GKE 管理的 GPU 设备插件 Daemonset。

应用 GKE GPU 驱动程序安装程序 DaemonSet 清单。此清单会在每个 A4X 节点上部署一个 GPU 驱动程序安装程序 Pod：

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/refs/heads/master/nvidia-driver-installer/cos/daemonset-preloaded.yaml

在 Google Kubernetes Engine (GKE) 上使用 NVIDIA GPU Operator 管理 GPU 栈：

在创建和设置 GPU 节点池部分中，按照从获取身份验证凭据这一步开始的说明操作。

安装 NVIDIA GPU Operator。完成所有步骤，但替换参考部分中用于使用 Helm 安装 NVIDIA GPU Operator 的命令。请改用以下命令：

helm install --wait --generate-name \
  -n gpu-operator \
  nvidia/gpu-operator \
  --version="25.3.0" \
  -f <(cat <<EOF
hostPaths:
  driverInstallDir: /home/kubernetes/bin/nvidia
toolkit:
  installDir: /home/kubernetes/bin/nvidia
cdi:
  enabled: true
  default: true
driver:
  enabled: false

daemonsets:
  tolerations:
    - key: nvidia.com/gpu
      operator: Equal
      value: present
      effect: NoSchedule
    - key: kubernetes.io/arch
      operator: Equal
      value: arm64
      effect: NoSchedule

node-feature-discovery:
  worker:
    tolerations:
      - key: kubernetes.io/arch
        operator: Equal
        value: arm64
        effect: NoSchedule
      - key: "node-role.kubernetes.io/master"
        operator: "Equal"
        value: ""
        effect: "NoSchedule"
      - key: "node-role.kubernetes.io/control-plane"
        operator: "Equal"
        value: ""
        effect: "NoSchedule"
      - key: nvidia.com/gpu
        operator: Exists
        effect: NoSchedule
EOF
)

安装 RDMA 二进制文件并配置 NCCL

应用以下 DaemonSet 以在每个节点上安装 RDMA 二进制文件和 NCCL 库。在每个底层虚拟机上，RDMA 二进制文件安装在 /home/kubernetes/bin/gib 目录中，NCCL 库安装在 /home/kubernetes/bin/nvidia/lib64 目录中。

  kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer-a4x.yaml

安装 NVIDIA Compute Domain CRD 和 DRA 驱动程序

安装 NVIDIA Compute Domain CRD 和 DRA 驱动程序。如需了解详情，请参阅 NVIDIA DRA GPU 驱动程序。

验证您已在开发环境中安装 Helm。Helm 已预安装在 Cloud Shell 中。

虽然没有特定的 Helm 版本要求，但您可以使用以下命令来验证 Helm 已安装。
```
helm version
```
如果输出类似于 Command helm not found，则可以通过运行以下命令来安装 Helm CLI：
```
curl -fsSL -o get_helm.sh https://raw.githubusercontent.com/helm/helm/master/scripts/get-helm-3 \
  && chmod 700 get_helm.sh \
  && ./get_helm.sh
```

添加 NVIDIA Helm 库：

helm repo add nvidia https://helm.ngc.nvidia.com/nvidia \
  && helm repo update

为 DRA 驱动程序创建 ResourceQuota：

export POD_QUOTA=POD_QUOTA

kubectl create ns nvidia-dra-driver-gpu

kubectl apply -n nvidia-dra-driver-gpu -f - << EOF
apiVersion: v1
kind: ResourceQuota
metadata:
  name: nvidia-dra-driver-gpu-quota
spec:
  hard:
    pods: ${POD_QUOTA}
  scopeSelector:
    matchExpressions:
    - operator: In
      scopeName: PriorityClass
      values:
        - system-node-critical
        - system-cluster-critical
EOF

将 POD_QUOTA 替换为至少是集群中 A4X 节点数量 2 倍的数字加 1。例如，如果集群中有 18 个 A4X 节点，则必须将该变量设置为至少 37。

安装 DRA 驱动程序：

helm install nvidia-dra-driver-gpu nvidia/nvidia-dra-driver-gpu \
    --version="DRA_DRIVER_VERSION" \
    --create-namespace \
    --namespace nvidia-dra-driver-gpu \
    -f <(cat <<EOF
nvidiaDriverRoot: /home/kubernetes/bin/nvidia
resources:
  gpus:
    enabled: false

controller:
  affinity:
      nodeAffinity:
        requiredDuringSchedulingIgnoredDuringExecution:
          nodeSelectorTerms:
          - matchExpressions:
            - key: "nvidia.com/gpu"
              operator: "DoesNotExist"

kubeletPlugin:
  affinity:
    nodeAffinity:
      requiredDuringSchedulingIgnoredDuringExecution:
        nodeSelectorTerms:
          - matchExpressions:
              - key: cloud.google.com/gke-accelerator
                operator: In
                values:
                  - nvidia-gb200
              - key: kubernetes.io/arch
                operator: In
                values:
                  - arm64

  tolerations:
    - key: nvidia.com/gpu
      operator: Equal
      value: present
      effect: NoSchedule
    - key: kubernetes.io/arch
      operator: Equal
      value: arm64
      effect: NoSchedule
EOF
)

将 DRA_DRIVER_VERSION 替换为 25.3.1 版或更高版本。您可以运行 helm search repo nvidia/nvidia-dra-driver-gpu --versions 命令来查找可用的 NVIDIA DRA 驱动程序版本列表。或者，查看 GitHub 版本说明，找到可用的版本。

为多网络、RDMA 和 IMEX 网域配置工作负载清单

将以下注解添加到 Pod 元数据中：

metadata:
  annotations:
    networking.gke.io/default-interface: 'eth0'
    networking.gke.io/interfaces: |
      [
        {"interfaceName":"eth0","network":"default"},
        {"interfaceName":"eth2","network":"rdma-0"},
        {"interfaceName":"eth3","network":"rdma-1"},
        {"interfaceName":"eth4","network":"rdma-2"},
        {"interfaceName":"eth5","network":"rdma-3"}
      ]

添加节点亲和性规则以在 Arm 节点上进行调度：

spec:
affinity:
  nodeAffinity:
    requiredDuringSchedulingIgnoredDuringExecution:
      nodeSelectorTerms:
      - matchExpressions:
        - key: kubernetes.io/arch
          operator: In
          values:
          - arm64

如需了解详情，请参阅将工作负载调度到单个架构。

将以下卷添加到 Pod 规范中：

spec:
volumes:
  - name: library-dir-host
    hostPath:
      path: /home/kubernetes/bin/nvidia
  - name: gib
    hostPath:
      path: /home/kubernetes/bin/gib

将以下卷装载、环境变量和资源添加到请求 GPU 的容器。工作负载容器必须请求所有四个 GPU：

containers:
- name: my-container
  volumeMounts:
    - name: library-dir-host
      mountPath: /usr/local/nvidia
    - name: gib
   mountPath: /usr/local/gib
  env:
    - name: LD_LIBRARY_PATH
      value: /usr/local/nvidia/lib64
  resources:
    limits:
      nvidia.com/gpu: 4

为工作负载创建 ComputeDomain 资源：

apiVersion: resource.nvidia.com/v1beta1
kind: ComputeDomain
metadata:
  name: a4x-compute-domain
spec:
  numNodes: NUM_NODES
  channel:
    resourceClaimTemplate:
      name: a4x-compute-domain-channel

将 NUM_NODES 替换为工作负载所需的节点数。

指定 Pod 将使用的 resourceClaimTemplate：

spec:
  ...
  volumes:
    ...
  containers:
    - name: my-container
      ...
      resources:
        limits:
          nvidia.com/gpu: 4
  claims:
          - name: compute-domain-channel
        ...
resourceClaims:
  - name: compute-domain-channel
    resourceClaimTemplateName: a4x-compute-domain-channel

设置所有必需的环境变量以配置 NCCL。使用工作负载容器中的以下 shell 脚本：
```
source /usr/local/gib/scripts/set_nccl_env.sh
```

完成后的 Pod 规范如下所示：

apiVersion: resource.nvidia.com/v1beta1
kind: ComputeDomain
metadata:
  name: a4x-compute-domain
spec:
  numNodes: NUM_NODES
  channel:
    resourceClaimTemplate:
      name: a4x-compute-domain-channel
---
apiVersion: apps/v1
kind: Pod
metadata:
  name: my-pod
  labels:
    k8s-app: my-pod
  annotations:
    networking.gke.io/default-interface: 'eth0'
    networking.gke.io/interfaces: |
      [
        {"interfaceName":"eth0","network":"default"},
        {"interfaceName":"eth2","network":"rdma-0"},
        {"interfaceName":"eth3","network":"rdma-1"},
        {"interfaceName":"eth4","network":"rdma-2"},
        {"interfaceName":"eth5","network":"rdma-3"},
      ]
spec:
  ...
  affinity:
    nodeAffinity:
      requiredDuringSchedulingIgnoredDuringExecution:
        nodeSelectorTerms:
        - matchExpressions:
          - key: kubernetes.io/arch
            operator: In
            values:
            - arm64
  volumes:
    - name: library-dir-host
      hostPath:
        path: /home/kubernetes/bin/nvidia
    - name: gib
      hostPath:
        path: /home/kubernetes/bin/gib
  containers:
    - name: my-container
      volumeMounts:
        - name: library-dir-host
          mountPath: /usr/local/nvidia
        - name: gib
       mountPath: /usr/local/gib
      env:
        - name: LD_LIBRARY_PATH
          value: /usr/local/nvidia/lib64
      resources:
        limits:
          nvidia.com/gpu: 4
        claims:
          - name: compute-domain-channel
        ...
resourceClaims:
  - name: compute-domain-channel
    resourceClaimTemplateName: a4x-compute-domain-channel

测试网络性能

我们建议您验证已配置集群的功能。为此，请使用 NCCL/gIB 测试，这些测试是针对 Google 环境优化的 NVIDIA Collective Communications Library (NCCL) 测试。

后续步骤

如需了解如何使用 TAS 和 Kueue 在 GKE 集群上调度工作负载，请参阅通过拓扑感知调度安排 GKE 工作负载。
如需了解如何管理与 GKE 集群和 AI 工作负载相关的常见事件，请参阅管理 AI 优化型 GKE 集群。