自動 NCCL インストーラを使用して GPUDirect RDMA を構成する

分散 AI / ML ワークロードを実行する場合、アプリケーションでは、複数のノードにまたがる GPU 間の高スループットで低レイテンシの通信が必要になることがよくあります。A4 仮想マシン（VM）インスタンスと A3 Ultra VM インスタンスには、NVIDIA ConnectX-7（CX7）ネットワーク インターフェース コントローラ（NIC）を基盤とする Titanium ML ネットワーク アダプタが搭載されています。どちらの VM ファミリーも、RDMA over Converged Ethernet（RoCE）を使用して、ノード間の GPU 間トラフィックを 3.2 Tbps の非ブロッキングで実現します。RoCE は、AI ワークロードに高性能なクラウド エクスペリエンスを提供することで、複数の GPU 間でのスケーリングとコラボレーションを可能にします。

GPUDirect RDMA はこの直接通信パスを提供しますが、必要なネットワーク プラグインとリソース クレームを手動で構成すると、複雑でエラーが発生しやすくなります。

この運用上の課題を解決するには、自動化された RDMA NCCL インストーラと自動化されたネットワーキング、Google Kubernetes Engine（GKE）マネージド DRANET を組み合わせて、GKE 上の複数のホストにわたる AI ワークロードの GPUDirect RDMA 構成を簡素化します。自動 RDMA NCCL インストーラは、A3 Ultra または A4 マシンタイプを使用するノードプールに、必要な NCCL バイナリと gIB バイナリを自動的にデプロイします。これらの機能を組み合わせて使用すると、基盤となるネットワーク インフラストラクチャとリソース クレームを手動で管理するのではなく、AI ワークロードのデプロイに集中できます。

詳細については、GPUDirect RDMA を使用してクラスタを作成するをご覧ください。

利点

自動 RDMA NCCL インストーラには、次の利点があります。

  • RDMA NCCL ライブラリの自動インストール: Standard クラスタの GKE バージョン 1.34.8-gke.1278000 以降と Autopilot クラスタのバージョン 1.36.2-gke.1368000 以降では、ノードプールの作成時に 1 つのフラグを使用して、 Google Cloud で最適な GPU ネットワーキング パフォーマンスを実現するように設計された RDMA NCCL/gIB ライブラリの自動インストールを有効にできます。
  • RDMA ワークロードを実行するための構成の簡素化: 自動化された RDMA NCCL インストーラ、アクセラレータ VM の自動ネットワーキング、GKE マネージド DRANET を組み合わせることで、手動でネットワークを構成することなく、RDMA 対応の AI ワークロードをデプロイできます。
  • 互換性と安全なアップグレード: インストーラによって提供される RDMA NCCL バイナリは、各 GKE リリースの前にエンドツーエンドの認定テストに合格しています。このテストにより、RDMA ユーザー空間ライブラリが GKE で使用される特定のカーネル バージョンとの互換性を維持し、実行中のワークロードを中断することなくクラスタとノードプールを安全にアップグレードできます。

制限事項

自動 RDMA NCCL インストーラには次の制限があります。

  • 自動 RDMA NCCL インストーラは、バージョン 1.34.8-gke.1278000 以降の GKE Standard クラスタと、バージョン 1.36.2-gke.1368000 以降の Autopilot クラスタでサポートされています。
  • 自動 RDMA NCCL インストーラは、a3-ultragpu-8g マシンタイプと a4-highgpu-8g マシンタイプでのみ使用できます。
  • GKE マネージド DRANET を使用してワークロードを構成する場合は、GKE マネージド DRANET の考慮すべきポイントと制限事項をご覧ください。
  • 同じノードプールでマルチネットワーク API と DRANET の両方を使用することはできません。Pod のネットワーク アタッチメントには、1 つの方法を選択する必要があります。
  • GKE ノードでは Container-Optimized OS ノードイメージを使用する必要があります。Ubuntu と Windows のノードイメージはサポートされません。
  • GKE で GPUDirect RDMA を使用して A3 Ultra と A4 を実行するための要件をすべて満たしている必要があります。

始める前に

作業を始める前に、次のことを確認してください。

  • クラスタで、Standard クラスタの場合は GKE バージョン 1.34.8-gke.1278000 以降、Autopilot クラスタの場合はバージョン 1.36.2-gke.1368000 以降が使用されていることを確認します。
  • Google Cloud CLI バージョン 550.0.0 以降を使用していることを確認します。
  • すでに Autopilot クラスタまたは Standard クラスタが存在していることを確認する。新しいクラスタを作成するには、Autopilot クラスタを作成するをご覧ください。
  • クラスタで GKE Dataplane V2 が有効になっていることを確認します。（--enable-dataplane-v2 フラグを指定して新しいクラスタを作成するときに、GKE Dataplane V2 を有効にできます）。GKE Dataplane V2 は、新しいすべての Autopilot クラスタでデフォルトで有効になります。
  • 次のいずれかの方法で、クラスタとノードプールのネットワーキングが設定されていることを確認します。
  • 仮想マシン（VM）インスタンスの割り当てが 2 つ以上あることを確認します。a4-highgpu-8g または a3-ultragpu-8g マシンタイプを使用する場合、この最小割り当ては 16 個の GPU に相当します。

インストーラと自動ネットワーキングを使用してノードプールを作成する

マルチホスト AI ワークロードで GPUDirect RDMA を有効にするには、自動 RDMA NCCL インストーラ、自動ネットワーキング、GKE マネージド DRANET を組み合わせた GPU ノードプールを構成します。

次のオプションのいずれかを選択します。

gcloud（Standard）

Standard ノードプールで自動 RDMA NCCL インストーラを有効にするには、--gpudirect-strategy フラグの値として rdma を指定します。GKE は、ノード バージョンと互換性のある RDMA NCCL プラグインを自動的にデプロイします。

  1. 次の環境変数を設定します。

    export CLUSTER_NAME=CLUSTER_NAME
export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
export PROJECT_ID=PROJECT_ID
export NODE_POOL_NAME=NODE_POOL_NAME
export NODEPOOL_LOCATION=NODEPOOL_LOCATION
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export ACCELERATOR_COUNT=ACCELERATOR_COUNT
export DRIVER_VERSION=DRIVER_VERSION
export MACHINE_TYPE=MACHINE_TYPE
export NUM_NODES=NUM_NODES
export RESERVATION_PROJECT=RESERVATION_PROJECT
export RESERVATION_NAME=RESERVATION_NAME
export RESERVATION_BLOCK=RESERVATION_BLOCK

    次のように置き換えます。

    • CLUSTER_NAME: クラスタの名前。
    • CONTROL_PLANE_LOCATION: クラスタ コントロール プレーンのロケーション。
    • PROJECT_ID: プロジェクト ID。
    • NODE_POOL_NAME: 新しいノードプールの名前。
    • NODEPOOL_LOCATION: ノードプールのロケーション。
    • ACCELERATOR_TYPE: アクセラレータ タイプ（nvidia-h200-141gb、nvidia-b200 など）。
    • ACCELERATOR_COUNT: アクセラレータの数（8 など）。
    • DRIVER_VERSION: GPU ドライバのバージョン（default、latest など）。
    • MACHINE_TYPE: マシンタイプ。a3-ultragpu-8g または a4-highgpu-8g にする必要があります。
    • NUM_NODES: ノードプール内のノード数。
    • RESERVATION_PROJECT: 予約が配置されているプロジェクト ID。
    • RESERVATION_NAME: 予約の名前。
    • RESERVATION_BLOCK: 予約ブロック。

  2. ノードプールを作成します。

    gcloud container node-pools create ${NODE_POOL_NAME} \
  --location=${CONTROL_PLANE_LOCATION} \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT_ID} \
  --node-locations=${NODEPOOL_LOCATION} \
  --accelerator type=${ACCELERATOR_TYPE},count=${ACCELERATOR_COUNT},gpu-driver-version=${DRIVER_VERSION} \
  --machine-type=${MACHINE_TYPE} \
  --num-nodes=${NUM_NODES} \
  --reservation-affinity=specific \
  --reservation=projects/${RESERVATION_PROJECT}/reservations/${RESERVATION_NAME}/reservationBlocks/${RESERVATION_BLOCK} \
  --accelerator-network-profile=auto \
  --node-labels=cloud.google.com/gke-networking-dra-driver=true \
  --gpudirect-strategy=rdma

ComputeClass（Autopilot と Standard）

Autopilot クラスタまたはノードプールの自動作成を使用する Standard クラスタで自動 RDMA NCCL インストーラを使用するには、gpuDirect: rdma、acceleratorNetworkProfile: auto、DRANET を有効にしてカスタム ComputeClass リソースを構成します。

  1. 次のマニフェストを rdma-compute-class.yaml として保存します。

    apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: COMPUTE_CLASS_NAME
spec:
  nodePoolConfig:
    dra:
      networking:
        enabled: true
  priorities:
  - gpu:
      type: ACCELERATOR_TYPE
      count: ACCELERATOR_COUNT
    machineType: MACHINE_TYPE
    # Optional: include the reservations block if you use a reservation
    reservations:
      affinity: Specific
      specific:
      - name: RESERVATION_NAME
        zones:
        - RESERVATION_ZONE
    acceleratorNetworkProfile: auto
    gpuDirect: rdma
  nodePoolAutoCreation:
    enabled: true

    次のように置き換えます。

    • COMPUTE_CLASS_NAME: ComputeClass の名前。
    • ACCELERATOR_TYPE: アクセラレータ タイプ（nvidia-h200-141gb、nvidia-b200 など）。
    • ACCELERATOR_COUNT: アクセラレータの数（8 など）。
    • MACHINE_TYPE: マシンタイプ。a3-ultragpu-8g または a4-highgpu-8g にする必要があります。
    • RESERVATION_NAME: 予約の名前。
    • RESERVATION_ZONE: 予約が配置されているゾーン。

  2. 次のようにマニフェストを適用します。

    kubectl apply -f rdma-compute-class.yaml

    この ComputeClass を使用するワークロードをデプロイする場合は、Pod 仕様に cloud.google.com/compute-class: COMPUTE_CLASS_NAME ノードセレクタを追加します。

Terraform

Terraform を使用して、自動化された RDMA NCCL インストーラが構成された GKE クラスタとノードプールをプロビジョニングできます。

  1. Terraform の公式ドキュメントに沿って、ローカルマシンに Terraform をインストールします。

  2. 構成を格納する空のディレクトリを作成し、そのディレクトリに移動します。main.tf という名前のファイルを作成し、次の内容を入力します。

    # main.tf

terraform {
  required_providers {
    google = {
      source  = "hashicorp/google"
      version = ">= 7.34.0" # Min version that supports the automated RDMA NCCL installer
    }
  }
}

provider "google" {
  project = "PROJECT_ID"
  region  = "REGION"
  zone    = "ZONE"
}

    次のように置き換えます。

    • PROJECT_ID: プロジェクト ID。
    • REGION: コンピューティング リージョン。
    • ZONE: A3 Ultra または A4 マシンタイプをサポートするゾーン。

  3. 同じディレクトリで次のコマンドを実行して、Google Cloud で認証し、Terraform の状態を初期化します。

    gcloud auth application-default login
terraform init

  4. main.tf または別の Terraform 構成ファイルに次の構成を追加します。この構成では、自動 RDMA NCCL インストーラ、自動アクセラレータ ネットワーク プロファイル、GKE マネージド DRANET を使用するように構成された 3 つの a3-ultragpu-8g ノードを含むクラスタを us-central1-b にデプロイします。

    resource "google_container_cluster" "cluster" {
  name               = "example-rdma-tf"
  location           = "us-central1-b"

  initial_node_count  = 3 # Number of CPU nodes in default pool
  deletion_protection = false # Optional, allows Terraform to delete clusters

  min_master_version      = "1.34.8-gke.1278000" # Required, min version for Standard clusters with automated RDMA NCCL installer support
  datapath_provider       = "ADVANCED_DATAPATH" # Required, enable GKE Dataplane v2
}

resource "google_container_node_pool" "rdma_pool" {
  name       = "rdma-pool"
  cluster    = google_container_cluster.cluster.name
  location   = google_container_cluster.cluster.location
  node_count = 3

  network_config {
    accelerator_network_profile = "auto" # Required if per-NIC networks are not manually configured
  }

  node_config {
    machine_type = "a3-ultragpu-8g"
    guest_accelerator {
      type  = "nvidia-h200-141gb"
      count = 8
    }

    ephemeral_storage_local_ssd_config { # Required, see https://github.com/hashicorp/terraform-provider-google/issues/17068
      data_cache_count = 0
      local_ssd_count = 32
    }

    gpudirect_strategy = "RDMA"

    labels = {
      "cloud.google.com/gke-networking-dra-driver"  = "true" # Automatically install DRANET driver
      "gke.networks.io/accelerator-network-profile" = "auto"
    }
  }
}

  5. terraform plan -out=tfplan を実行して、現在の構成に基づいて Terraform が行う変更を確認します。プランを確認したら、次のコマンドを使用して実行します。

    terraform apply tfplan

ResourceClaimTemplate を作成する

Pod が RDMA デバイスを動的にリクエストして受信できるようにするには、ResourceClaimTemplate オブジェクトを作成します。これにより、AI ワークロードに高性能なネットワーキングを提供できます。allocationMode: All 設定では、使用可能なすべての mrdma.google.com RDMA デバイスが割り当てられます。使用可能なすべてのデバイスを割り当てることは、GPU と RDMA の機能を最大限に活用する必要がある複数のホストにまたがる AI ワークロードにとって重要です。

  1. 次のマニフェストを all-mrdma-template.yaml として保存します。

    apiVersion: resource.k8s.io/v1
kind: ResourceClaimTemplate
metadata:
  name: all-mrdma
spec:
  spec:
    devices:
      requests:
      - name: req-mrdma
        exactly:
          deviceClassName: mrdma.google.com
          allocationMode: All

  2. 次のようにマニフェストを適用します。

    kubectl apply -f all-mrdma-template.yaml

NCCL RDMA テスト ワークロードをデプロイして実行する

RDMA 構成を確認するには、テスト ワークロードをデプロイし、NCCL all-gather テストを実行して設定を確認します。

  1. 次のマニフェストを nccl-test-workload.yaml として保存します。

    apiVersion: v1
kind: Service
metadata:
  name: nccl-host-1
spec:
  selector:
    name: nccl-host-1
  clusterIP: None
---
apiVersion: v1
kind: Service
metadata:
  name: nccl-host-2
spec:
  selector:
    name: nccl-host-2
  clusterIP: None
---
apiVersion: v1
kind: Pod
metadata:
  name: nccl-test-host-1
  labels:
    name: nccl-host-1
spec:
  # Uncomment the following lines if you use a custom ComputeClass:
  # nodeSelector:
  #   cloud.google.com/compute-class: COMPUTE_CLASS_NAME
  volumes:
    - name: library-dir-host
      hostPath:
        path: /home/kubernetes/bin/nvidia
    - name: gib
      hostPath:
        path: /home/kubernetes/bin/gib
    - name: shared-memory
      emptyDir:
        medium: "Memory"
        sizeLimit: 250Gi
  resourceClaims:
  - name: rdma
    resourceClaimTemplateName: all-mrdma
  containers:
    - image: us-docker.pkg.dev/gce-ai-infra/gpudirect-gib/nccl-plugin-gib-diagnostic:v1.1.2
      name: test
      resources:
        requests:
          cpu: 150m
        limits:
          nvidia.com/gpu: 8
        claims:
        - name: rdma
      volumeMounts:
        - name: library-dir-host
          mountPath: /usr/local/nvidia
        - name: gib
          mountPath: /usr/local/gib
        - name: shared-memory
          mountPath: /dev/shm
      env:
        - name: LD_LIBRARY_PATH
          value: /usr/local/nvidia/lib64
      command: ["/bin/bash", "-c"]
      args:
        - |
          /scripts/container_entry.sh shell
          source /usr/local/gib/scripts/set_nccl_env.sh
          sleep infinity
---
apiVersion: v1
kind: Pod
metadata:
  name: nccl-test-host-2
  labels:
    name: nccl-host-2
spec:
  # Uncomment the following lines if you use a custom ComputeClass:
  # nodeSelector:
  #   cloud.google.com/compute-class: COMPUTE_CLASS_NAME
  volumes:
    - name: library-dir-host
      hostPath:
        path: /home/kubernetes/bin/nvidia
    - name: gib
      hostPath:
        path: /home/kubernetes/bin/gib
    - name: shared-memory
      emptyDir:
        medium: "Memory"
        sizeLimit: 250Gi
  resourceClaims:
  - name: rdma
    resourceClaimTemplateName: all-mrdma
  containers:
    - image: us-docker.pkg.dev/gce-ai-infra/gpudirect-gib/nccl-plugin-gib-diagnostic:v1.1.2
      name: test
      resources:
        requests:
          cpu: 150m
        limits:
          nvidia.com/gpu: 8
        claims:
        - name: rdma
      volumeMounts:
        - name: library-dir-host
          mountPath: /usr/local/nvidia
        - name: gib
          mountPath: /usr/local/gib
        - name: shared-memory
          mountPath: /dev/shm
      env:
        - name: LD_LIBRARY_PATH
          value: /usr/local/nvidia/lib64
      command: ["/bin/bash", "-c"]
      args:
        - |
          /scripts/container_entry.sh shell
          source /usr/local/gib/scripts/set_nccl_env.sh
          sleep infinity

  2. テスト ワークロード マニフェストを適用します。

    kubectl apply -f nccl-test-workload.yaml

  3. 2 つのテストホスト間で NCCL all-gather テストを実行します。

    kubectl exec nccl-test-host-1 -it -- /usr/local/gib/scripts/run_nccl_tests.sh -t all_gather -b 1K -e 8G nccl-host-1 nccl-host-2

AI / ML ワークロードをデプロイする

分散トレーニング ジョブを効率的に実行するには、ML アプリケーションがインストーラによってプロビジョニングされた RDMA デバイスをリクエストする必要があります。テスト ワークロードで RDMA 構成を確認したら、AI / ML アプリケーションをデプロイできます。アプリケーション Pod は、以前に作成した ResourceClaimTemplate オブジェクト（all-mrdma）を参照して、RDMA デバイスを動的にリクエストして受信する必要があります。

次の Pod マニフェストでは、all-mrdma ResourceClaimTemplate オブジェクトを使用します。IMAGE は AI / ML ワークロード コンテナ イメージに、COMMAND はアプリケーションのエントリ ポイントに置き換えます。

apiVersion: v1
kind: Pod
metadata:
  name: my-rdma-workload
spec:
  restartPolicy: OnFailure
  # Uncomment the following lines if you use a custom ComputeClass:
  # nodeSelector:
  #   cloud.google.com/compute-class: COMPUTE_CLASS_NAME
  volumes:
    - name: library-dir-host
      hostPath:
        path: /home/kubernetes/bin/nvidia
    - name: gib
      hostPath:
        path: /home/kubernetes/bin/gib
    - name: shared-memory
      emptyDir:
        medium: "Memory"
        sizeLimit: 250Gi
  resourceClaims:
  - name: mrdma-claim
    resourceClaimTemplateName: all-mrdma
  containers:
  - name: my-container
    image: IMAGE # Replace with your AI/ML workload image, for example, gcr.io/your-project/your-ai-app:latest
    command: ["/bin/bash", "-c"]
    args:
      - |
        source /usr/local/gib/scripts/set_nccl_env.sh
        COMMAND # Replace with your workload command, for example, python3 train.py
    resources:
      limits:
        nvidia.com/gpu: 8 # Specify the required GPU count
      claims:
      - name: mrdma-claim
    volumeMounts:
      - name: library-dir-host
        mountPath: /usr/local/nvidia
      - name: gib
        mountPath: /usr/local/gib
      - name: shared-memory
        mountPath: /dev/shm
    env:
      - name: LD_LIBRARY_PATH
        value: /usr/local/nvidia/lib64
      - name: MASTER_ADDR
        value: "nccl-test-host-1" # Example for distributed training
      - name: MASTER_PORT
        value: "29500"

このワークロードをデプロイするには、マニフェストを workload.yaml として保存し、適用します。

kubectl apply -f workload.yaml

高度な構成

DRANET による自動ネットワーキングではなく、手動のマルチネットワーク構成を使用するノードプールで、自動 RDMA NCCL インストーラを有効にすることもできます。

マルチネットワークが有効になっているクラスタにノードプールを作成する

GKE での GPUDirect RDMA のサポートのクラスタ構成に沿って操作した場合は、次の手順で、マルチネットワークが有効になっている既存のクラスタで自動 RDMA NCCL インストーラを有効にして、新しいノードプールを作成できます。

クラスタでマルチネットワークが有効になっているかどうかを確認するには、次のコマンドを実行します。

export CLUSTER_NAME=CLUSTER_NAME
export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
export PROJECT_ID=PROJECT_ID
export CLUSTER_VERSION=CLUSTER_VERSION

gcloud container clusters describe ${CLUSTER_NAME} \
  --location ${CONTROL_PLANE_LOCATION} --project ${PROJECT_ID} \
  --format="value(networkConfig.enableMultiNetworking)"

A3 Ultra VM インスタンスと A4 VM インスタンスの追加 NIC に VPC ネットワークを再利用するには、次の変数を設定します。

export REGION=REGION
export ZONE=ZONE
export GVNIC_NETWORK_PREFIX=GVNIC_NETWORK_PREFIX
export RDMA_NETWORK_PREFIX=RDMA_NETWORK_PREFIX

自動 RDMA NCCL インストーラを使用してノードプールを作成するには、次のコマンドを実行します。

export NODE_POOL_NAME=NODE_POOL_NAME
export PROJECT_ID=PROJECT_ID
export NODE_COUNT=NODE_COUNT
export ACCELERATOR_TYPE=ACCELERATOR_TYPE
export ACCELERATOR_COUNT=ACCELERATOR_COUNT
export DRIVER_VERSION=DRIVER_VERSION
export MACHINE_TYPE=MACHINE_TYPE
export RESERVATION=RESERVATION

gcloud container node-pools create ${NODE_POOL_NAME} --project ${PROJECT_ID} \
  --location ${CONTROL_PLANE_LOCATION} --node-locations ${ZONE} --cluster ${CLUSTER_NAME} --num-nodes=${NODE_COUNT} \
  --machine-type ${MACHINE_TYPE} \
  --accelerator type=${ACCELERATOR_TYPE},count=${ACCELERATOR_COUNT},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 \
  --additional-node-network network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-4 \
  --additional-node-network network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-5 \
  --additional-node-network network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-6 \
  --additional-node-network network=${RDMA_NETWORK_PREFIX}-net,subnetwork=${RDMA_NETWORK_PREFIX}-sub-7 \
  --scopes "https://www.googleapis.com/auth/cloud-platform" \
  --reservation-affinity=specific \
  --reservation=${RESERVATION} \
  --gpudirect-strategy=rdma

次のステップ