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

本文档介绍了如何创建使用 A4X Max Compute Engine 实例的 AI 优化型 Google Kubernetes Engine (GKE) 集群,以支持您的 AI 和机器学习工作负载。

借助 A4X Max 和 A4X 系列,您可以使用 NVIDIA 多节点 NVLink (MNNVL) 系统(一种机架级解决方案,可提供更高的 GPU 功率和性能)运行 大规模 AI/机器学习集群。这些机器提供有针对性的工作负载放置、拓扑感知调度和高级集群维护控制等功能。如需了解详情,请参阅 集群管理功能。 借助 A4X Max,GKE 还提供自动网络设置,可简化集群配置。

AI 和机器学习工作负载(例如分布式训练)需要通过缩短作业完成时间来获得极大的加速,从而优化性能。GKE 提供了一个统一的平台界面,可为您的组织运行各种工作负载,从而减轻管理多个平台的运营负担。 您可以运行高性能分布式预训练、模型微调、模型推理、应用服务和支持服务等工作负载。对于需要高性能、高吞吐量和低延迟的工作负载,GPUDirect RDMA 可减少将载荷传输到 GPU 和从 GPU 传输载荷所需的网络跃点数。这种方法可以更高效地利用可用的网络带宽。如需了解详情,请参阅 GPU 网络 堆栈

在本文档中,您将了解如何使用 Google Cloud CLI 创建 GKE 集群,以便根据工作负载的需求灵活地配置集群。如需使用 gcloud CLI 创建具有其他机器类型的集群,请参阅以下内容:

或者,您也可以选择使用 Cluster Toolkit 快速部署集群,其中包含反映许多使用场景的最佳实践的默认设置。如需了解详情,请参阅创建具有默认配置的 AI 优化型 GKE集群。

准备工作

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

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

获取容量

您可以通过创建未来预留来获取 A4X Max 计算实例的容量。如需详细了解未来预留,请参阅Future reservations in AI Hypercomputer表格中的 Choose a consumption option列。

如需通过未来预留获取容量,请参阅AI Hypercomputer 中的 未来预留 行,该行位于如何获取 容量表格中。

要求

以下要求适用于具有 A4X Max 计算实例的 AI 优化型 GKE 集群:

  • 对于 A4X Max,您必须使用以下版本之一:

    • 对于 1.35 或更高版本,请使用 GKE 1.35.0-gke.2745000 或更高版本。
    • 对于 1.34,请使用 GKE 1.34.3-gke.1318000 或更高版本。

    这些版本有助于确保 A4X Max 使用以下内容:

    • R580.95.05,A4X Max 的最低 GPU 驱动程序版本,默认处于启用状态。
    • 一致的基于驱动程序的内存管理 (CDMM),默认处于启用状态。NVIDIA 建议 Kubernetes 集群启用此模式,以解决内存报告过多的问题。CDMM 允许通过驱动程序(而不是操作系统 (OS))管理 GPU 内存。这种方法有助于避免 OS 在线处理 GPU 内存,并将 GPU 内存作为非统一内存访问 (NUMA) 节点公开给 OS。启用 CDMM 后,不支持多实例 GPU。如需详细了解 CDMM,请参阅硬件和软件 支持
    • GPUDirect RDMA 和 MNNVL,建议启用这些功能,以便 A4X Max 节点池使用 A4X Max 的网络功能。
  • GKE 节点必须使用 Container-Optimized OS 节点映像。不支持 Ubuntu 和 Windows 节点映像。

  • 您的 GKE 工作负载必须使用单个 GKE 节点上的所有可用 GPU,并且您的 Pod 必须使用单个 GKE 节点上的所有可用次要 NIC。 多个 Pod 无法在单个 GKE 节点上共享 RDMA。

  • 您必须使用受预留约束的 预配模型来创建具有 A4X Max 的集群。不支持其他预配 模型。

  • 以下说明使用 DRANET 配置具有 A4X Max 的 AI 优化型 GKE 集群。 a4x-maxgpu-4g-metal 机器类型不支持多网络。

创建集群的注意事项

创建集群时,请考虑以下信息:

  • 选择集群位置
    • 验证您使用的位置是否提供您选择的机器类型。如需了解详情,请参阅加速器 可用性
    • 区域级 集群(建议用于生产工作负载)中创建节点池时,您可以使用 --node-locations 标志为 GKE 节点指定可用区。
  • 选择驱动程序版本
    • 驱动程序版本可以是以下某个值:
      • default:为您的 GKE 节点版本安装默认驱动程序版本。如需详细了解 默认驱动程序版本的要求,请参阅 要求 部分。
      • latest:为您的 GKE 版本安装最新可用的驱动程序版本。此选项仅适用于使用 Container-Optimized OS 的节点。
      • disabled:跳过自动驱动程序安装。您必须手动 安装 驱动程序 后创建节点池。
    • 如需详细了解 GKE 节点版本的默认和最新 GPU 驱动程序版本 ,请参阅 手动安装 NVIDIA GPU 驱动程序部分中的表格。
  • 选择预留亲和性

    • 您可以找到有关预留的信息,例如预留的名称或预留中特定块的名称。如需 查找这些值,请参阅查看未来预留 请求
    • --reservation-affinity 标志可以采用 specificany 值。不过,对于高性能分布式 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 替换为子块的名称。

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

对于分布式 AI 工作负载,通常会将多个 GPU 节点链接在一起,以作为一台计算机运行。A4X Max 是基于 NVIDIA GB300 NVL72 机架级架构的艾级平台。A4X Max 计算实例使用多层级分层网络架构,并采用导轨对齐设计,以优化各种通信类型的性能。此机器类型通过为 AI 工作负载提供高性能云体验,实现跨多个 GPU 的伸缩和协作。如需详细了解 A4X Max 的网络架构 (包括网络带宽和 NIC 配置),请参阅 A4X Max 机器类型(裸 金属)

如需创建使用 GPUDirect RDMA 和 MNNVL 的具有 A4X Max 的 GKE Standard 集群,请完成以下部分中所述的步骤:

  1. 创建 GKE 集群
  2. 创建工作负载政策
  3. 创建具有 A4X Max 的节点池
  4. 使用 asapd-lite 配置 MRDMA NIC
  5. 安装 NVIDIA Compute Domain CRD 和 DRA 驱动程序
  6. 为 RDMA 和 IMEX 网域配置工作负载清单

以下说明使用加速器网络 配置文件自动为 A4X Max 节点配置 VPC 网络和子网。 或者,您也可以明确指定 VPC 网络和 子网

创建 GKE 集群

  1. 创建 GKE Standard 集群:

    gcloud container clusters create CLUSTER_NAME \
      --enable-dataplane-v2 \
      --enable-ip-alias \
      --location=COMPUTE_REGION \
      --cluster-version=CLUSTER_VERSION \
      --no-enable-shielded-nodes [\
      --services-ipv4-cidr=SERVICE_CIDR \
      --cluster-ipv4-cidr=POD_CIDR \
      --addons=GcpFilestoreCsiDriver=ENABLED]
    

    替换以下内容:

    • CLUSTER_NAME:您的集群的名称。
    • CLUSTER_VERSION:新集群的版本。 如需详细了解哪个版本的 GKE 支持 您的配置,请参阅本文档中的要求
    • COMPUTE_REGION:计算区域的名称。
    • (可选)您可以明确指定服务和 Pod 的次要 CIDR 范围。如果您使用这些可选标志,请替换以下变量:

      • SERVICE_CIDR:服务的次要 CIDR 范围。
      • POD_CIDR:Pod 的次要 CIDR 范围。

      使用这些标志时,您必须验证 CIDR 范围是否与额外节点网络的子网范围重叠。例如, 考虑 SERVICE_CIDR=10.65.0.0/19POD_CIDR=10.64.0.0/19。如需了解详情,请参阅添加 Pod IPv4 地址 范围

  2. 如需在后续部分中运行 kubectl 命令,请连接到您的集群:

    gcloud container clusters get-credentials CLUSTER_NAME --location=COMPUTE_REGION
    

    替换以下内容:

    • CLUSTER_NAME:您的集群的名称。
    • COMPUTE_REGION:计算区域的名称。

    如需了解详情,请参阅安装 kubectl 并配置集群 访问权限

创建工作负载政策

必须要有工作负载政策才能创建分区。如需了解详情,请参阅 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 Max 的节点池

  1. 创建以下配置文件,以便使用节点池预分配 巨页

    cat > node_custom.yaml <<EOF
    linuxConfig:
      hugepageConfig:
        hugepage_size2m: 4096
    EOF
    
    export NODE_CUSTOM=node_custom.yaml
    
  2. 创建 A4X Max 节点池:

    gcloud container node-pools create NODE_POOL_NAME \
        --cluster=CLUSTER_NAME \
        --location=COMPUTE_REGION \
        --node-locations=COMPUTE_ZONE \
        --num-nodes=NODE_COUNT \
        --placement-policy=WORKLOAD_POLICY_NAME \
        --machine-type=a4x-maxgpu-4g-metal \
        --accelerator=type=nvidia-gb300,count=4,gpu-driver-version=latest \
        --system-config-from-file=${NODE_CUSTOM} \
        --accelerator-network-profile=auto \
        --node-labels=cloud.google.com/gke-networking-dra-driver=true,cloud.google.com/gke-dpv2-unified-cni=cni-migration \
        --reservation-affinity=specific \
        --reservation=RESERVATION_NAME/reservationBlocks/BLOCK_NAME/reservationSubBlocks/SUB_BLOCK_NAME
    

    替换以下内容:

    • NODE_POOL_NAME:节点池的名称。
    • CLUSTER_NAME:您的集群的名称。
    • COMPUTE_REGION:集群的计算区域。
    • COMPUTE_ZONE:节点池的可用区。
    • NODE_COUNT:节点池的节点数,必须为 18 个或更少。我们建议使用 18 个节点,以便使用 NVLink 网域在一个子块中获取 1x72 的 GPU 拓扑。
    • WORKLOAD_POLICY_NAME:您之前创建的工作负载政策的名称。
    • RESERVATION_NAME:预留的名称。 如需查找此值,请参阅查看未来预留 请求
    • BLOCK_NAME:预留中特定块的名称。如需查找此值,请参阅查看未来预留 请求

    此命令使用 auto 加速器网络配置文件自动创建一个网络,该网络连接单个可用区内的所有 A4X Max 节点。 当您使用 --accelerator-network-profile=auto 标志创建节点池时,GKE 会自动向节点添加 gke.networks.io/accelerator-network-profile: auto 标签。如需在这些节点上调度工作负载,您必须在工作负载的 nodeSelector 字段中添加此标签。

使用 asapd-lite 配置 MRDMA NIC

asapd-lite DaemonSet 配置 MRDMA NIC。不健康的 asapd-lite DaemonSet 可能表示没有 RDMA 连接。

  1. 安装 DaemonSet:

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/refs/heads/master/asapd-lite-installer/asapd-lite-installer-a4x-max-bm-cos.yaml
    
  2. 验证 asapd-lite DaemonSet 中的副本:

    kubectl get daemonset -n kube-system asapd-lite
    

    输出类似于以下内容:

    NAME         DESIRED   CURRENT   READY   UP-TO-DATE   AVAILABLE   NODE SELECTOR   AGE
    asapd-lite   18        18        18      18           18          <none>          5m
    

    READY 副本的数量应与节点池中已创建且运行正常的节点数一致。

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

以下步骤将安装 NVIDIA Compute Domain CRD 和 DRA 驱动程序,以启用 MNNVL 的使用。如需了解详情,请参阅GPU 的 NVIDIA DRA 驱动程序

  1. 验证您已在开发环境中安装 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
    
  2. 添加 NVIDIA Helm 库:

    helm repo add nvidia https://helm.ngc.nvidia.com/nvidia \
      && helm repo update
    
  3. 为 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 Max 节点数 2 倍加 1 的数字。例如,如果集群中有 18 个 A4X Max 节点,则必须将变量设置为至少 37。

  4. 安装 ComputeDomain CRD 和 DRA 驱动程序:

    helm install nvidia-dra-driver-gpu nvidia/nvidia-dra-driver-gpu \
        --set controller.args.v=4 --set kubeletPlugin.args.v=4 \
        --version="25.8.0" \
        --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-gb300
                  - 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
    )
    

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

  1. 添加节点亲和性规则,以在 Arm 节点上调度工作负载:

    spec:
      affinity:
        nodeAffinity:
          requiredDuringSchedulingIgnoredDuringExecution:
            nodeSelectorTerms:
            -   matchExpressions:
              -   key: kubernetes.io/arch
                operator: In
                values:
                -   arm64
    
  2. 将以下卷添加到 Pod 规范:

    spec:
      volumes:
        - name: library-dir-host
          hostPath:
            path: /home/kubernetes/bin/nvidia
    
  3. 将以下卷装载、环境变量和资源添加到请求 GPU 的容器。您的工作负载容器必须请求所有四个 GPU:

    containers:
      - name: my-container
        volumeMounts:
          - name: library-dir-host
            mountPath: /usr/local/nvidia
    
        env:
          - name: LD_LIBRARY_PATH
            value: /usr/local/nvidia/lib64
        resources:
          limits:
            nvidia.com/gpu: 4
    
  4. 为工作负载创建 ComputeDomain 资源:

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

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

  5. 创建一个 ResourceClaimTemplate,以使用 DRANET 分配网络资源,并为 Pod 请求 RDMA 设备:

    apiVersion: resource.k8s.io/v1
    kind: ResourceClaimTemplate
    metadata:
      name: all-mrdma
    spec:
      spec:
        devices:
          requests:
          - name: req-mrdma
            exactly:
              deviceClassName: mrdma.google.com
              allocationMode: ExactCount
              count: 8
    
  6. 指定 Pod 使用的 ResourceClaimTemplate:

    spec:
      ...
      volumes:
        ...
      containers:
        - name: my-container
          ...
          resources:
            limits:
              nvidia.com/gpu: 4
      claims:
              - name: compute-domain-channel
              - name: rdma
            ...
    resourceClaims:
      - name: compute-domain-channel
        resourceClaimTemplateName: a4x-max-compute-domain-channel
      - name: rdma
        resourceClaimTemplateName: all-mrdma
    
  7. 在容器映像中安装最新的用户空间库:

    Debian 12+/Ubuntu 20.04+(.deb 软件包)

    apt update
    apt install curl
    
    # Fetch DOCA OFED userspace libraries
    export DOCA_URL="https://linux.mellanox.com/public/repo/doca/latest/ubuntu22.04/arm64-sbsa/"
    export BASE_URL="https://linux.mellanox.com/public/repo/doca"
    curl "${BASE_URL}/GPG-KEY-Mellanox.pub" | gpg --dearmor -o /etc/apt/trusted.gpg.d/GPG-KEY-Mellanox.pub
    echo "deb [signed-by=/etc/apt/trusted.gpg.d/GPG-KEY-Mellanox.pub] ${DOCA_URL} ./" | tee /etc/apt/sources.list.d/doca.list
    
    apt update
    apt install doca-ofed-userspace
    
    # Upgrade to latest NCCL for best compatibility
    apt install --only-upgrade --allow-change-held-packages libnccl2 libnccl-dev

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

apiVersion: resource.nvidia.com/v1beta1
kind: ComputeDomain
metadata:
  name: a4x-max-compute-domain
spec:
  numNodes:  NUM_NODES
  channel:
    resourceClaimTemplate:
      name: a4x-max-compute-domain-channel
---
apiVersion: apps/v1
kind: Pod
metadata:
  name: my-pod
  labels:
    k8s-app: my-pod
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
  hostNetwork: true
  containers:
    - name: my-container
      volumeMounts:
        - name: library-dir-host
          mountPath: /usr/local/nvidia
      env:
        - name: LD_LIBRARY_PATH
          value: /usr/local/nvidia/lib64
      resources:
        limits:
          nvidia.com/gpu: 4
        claims:
          - name: compute-domain-channel
          - name: rdma
        ...
  resourceClaims:
    - name: compute-domain-channel
      resourceClaimTemplateName: a4x-max-compute-domain-channel
    - name: rdma
      resourceClaimTemplateName: all-mrdma

测试网络性能

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

如需了解详情,请参阅在自定义 GKE 集群上运行 NCCL,这些集群 使用 A4X Max

后续步骤