建立使用 A4X Max 的 AI 適用 GKE 叢集

本文件說明如何建立 AI 最佳化 Google Kubernetes Engine (GKE) 叢集,並使用 A4X Max Compute Engine 執行個體支援 AI 和機器學習工作負載。

A4X Max 和 A4X 系列可讓您使用 NVIDIA 多節點 NVLink (MNNVL) 系統執行大規模 AI/機器學習叢集,這是一種機架規模解決方案,可提高 GPU 效能。這些機器提供多項功能,例如指定工作負載位置、拓撲感知調度,以及進階叢集維護控制項。詳情請參閱「叢集管理功能」。有了 A4X Max,GKE 還能自動設定網路,簡化叢集設定。

AI 和機器學習工作負載 (例如分散式訓練) 需要強大的加速功能,才能縮短工作完成時間,進而最佳化效能。GKE 提供單一平台介面,可為貴機構執行多種工作負載,減少管理多個平台的營運負擔。您可以執行高效能分散式預先訓練、模型微調、模型推論、應用程式服務和支援服務等工作負載。對於需要高效能、高處理量和低延遲的工作負載,GPUDirect RDMA 可減少在 GPU 之間傳輸酬載所需的網路躍點。這種做法能更有效率地運用可用的網路頻寬。詳情請參閱「GPU 網路堆疊」。

本文說明如何使用 Google Cloud CLI 建立 GKE 叢集,以便根據工作負載需求,靈活設定叢集。如要使用 gcloud CLI 建立其他機器類型的叢集,請參閱下列內容:

或者,您也可以選擇使用 Cluster Toolkit,以預設設定快速部署叢集,這些設定反映了許多用途的最佳做法。詳情請參閱「使用 Cluster Toolkit 建立 AI 最佳化 GKE 叢集」。

事前準備

開始之前,請確認您已執行下列工作:

  • 啟用 Google Kubernetes Engine API。
  • 啟用 Google Kubernetes Engine API
  • 如要使用 Google Cloud CLI 執行這項工作,請安裝並初始化 gcloud CLI。如果您先前已安裝 gcloud CLI,請執行 gcloud components update 指令,取得最新版本。較舊的 gcloud CLI 版本可能不支援執行本文中的指令。

必要的角色

如要取得建立及管理 GKE 叢集所需的權限,請要求管理員授予您下列 IAM 角色:

如要進一步瞭解如何授予角色,請參閱「管理專案、資料夾和組織的存取權」。

您或許也能透過自訂角色或其他預先定義的角色,取得必要權限。

取得運算資源

您可以建立未來預留項目,取得 A4X Max 運算執行個體的容量。如要進一步瞭解未來預留項目,請參閱「選擇使用模式」表格中的「AI Hypercomputer 的未來預留項目」一欄。

如要透過未來預留項目取得容量,請參閱「如何取得容量」表格中的「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 可透過驅動程式 (而非作業系統) 管理 GPU 記憶體。這種做法可避免 OS 將 GPU 記憶體上線,並將 GPU 記憶體做為非統一記憶體存取 (NUMA) 節點,公開給 OS。啟用 CDMM 時,不支援多實體 GPU。如要進一步瞭解 CDMM,請參閱「硬體和軟體支援」。
    • GPUDirect RDMA 和 MNNVL:建議啟用,讓 A4X Max 節點集區使用 A4X Max 的網路功能。
  • GKE 節點必須使用 Container-Optimized OS 節點映像檔。不支援 Ubuntu 和 Windows 節點映像檔。

  • 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 旗標可採用 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 替換為子區塊的名稱。

建立 AI 最佳化 GKE 叢集,使用 A4X Max 和 GPUDirect RDMA

如果是分散式 AI 工作負載,多個 GPU 節點通常會連結在一起,做為單一電腦運作。A4X Max 是採用 NVIDIA GB300 NVL72 機架規模架構的百億億次級平台,A4X Max 運算執行個體採用多層級、階層式網路架構,並採用符合軌道 (rail-aligned) 設計,可針對各種通訊類型最佳化效能。這個機型可為 AI 工作負載提供高效能的雲端體驗,讓您在多個 GPU 之間擴充及協作。如要進一步瞭解 A4X Max 的網路架構,包括網路頻寬和 NIC 排列方式,請參閱「A4X Max 機型 (裸機)」。

如要建立採用 A4X Max 的 GKE Standard 叢集,並使用 GPUDirect RDMA 和 MNNVL,請完成下列各節所述的步驟:

  1. 建立 GKE 叢集
  2. 建立工作負載政策
  3. 建立 A4X Max 節點集區
  4. 使用 asapd-lite 設定 MRDMA NIC
  5. 安裝 NVIDIA Compute Domain CRD 和 DRA 驅動程式
  6. 設定 RDMA 和 IMEX 網域的工作負載資訊清單

這些操作說明會使用加速晶片 VM 網路設定檔,自動為 A4X Max 節點設定虛擬私有雲網路和子網路。或者,您也可以明確指定虛擬私有雲網路和子網路。

建立 GKE 叢集

  1. 建立 GKE Standard 叢集:

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

    更改下列內容:

    • CLUSTER_NAME:叢集名稱。
    • CLUSTER_VERSION:新叢集的版本。如要進一步瞭解哪些 GKE 版本支援您的設定,請參閱本文的「需求條件」一節。
    • COMPUTE_REGION:運算區域的名稱。
    • 您也可以視需要明確提供服務和 Pod 的次要 CIDR 範圍。如果您使用這些選用標記,請替換下列變數:

      • SERVICE_CIDR:Service 的次要 CIDR 範圍。
      • POD_CIDR:Pod 的次要 CIDR 範圍。

      使用這些旗標時,請務必確認 CIDR 範圍不會與其他節點網路的子網路範圍重疊。例如,請考慮 SERVICE_CIDR=10.65.0.0/19 和 POD_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. 建立下列設定檔,以便透過節點集區預先分配 hugepage:

    cat > node_custom.yaml <<EOF
    linuxConfig:
      hugepageConfig:
        hugepage_size2m: 4096
    EOF
    
    export NODE_CUSTOM=node_custom.yaml
    
  2. 建立 A4X Max 節點集區。--accelerator-network-profile=auto 旗標會自動建立網路,連結單一可用區內的所有 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:預留項目中特定區塊的名稱。如要查看這個值,請參閱「查看未來預留項目要求」。

使用 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。Cloud Shell 已預先安裝 Helm。

    雖然沒有特定的 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) 測試。

詳情請參閱「在採用 A4X Max 的 GKE 叢集上執行 NCCL」。

後續步驟