Integrator Guide

GKE GB200 (A4X) Networking Prerequisites

View as Markdown

For the GB200 GKE COS recipes (gb200-gke-cos-training, gb200-gke-cos-training-kubeflow, gb200-gke-cos-training-slurm, gb200-gke-cos-inference, and gb200-gke-cos-inference-dynamo, all on a4x-highgpu-4g nodes), GPUDirect-RDMA over RoCE enables high-speed inter-node GPU communication on GKE. AICR’s own GB200/GKE validation workload sets NCCL_NET=gIB explicitly (see recipes/components/gke-gb200-rdma/manifests/nccl-gib-installer-arm64.yaml and the runtime it applies, validators/performance/testdata/gb200/gke/runtime-nvls.yaml) rather than letting NCCL auto-select a plugin, so a missing or misconfigured RDMA fabric doesn’t silently fall back to a slower network path. It fails outright instead.

This fail-closed behavior belongs to the validation workload, not to every workload the recipe can run. For Kubeflow, the gb200-gke-cos-training-kubeflow leaf also ships that wiring pre-built, as the torch-distributed-rdma runtime below. A workload that uses the generic torch-distributed runtime instead, or authors its own TrainingRuntime, carries none of it.

GPUDirect RDMA on a4x-highgpu-4g is also incompatible with NCCL Fast Socket and the GPUDirect TCPX/TCPXO plugin (see GKE TCPXO Networking for that alternative, non-RDMA path); don’t enable either on a cluster that uses RDMA.

Infrastructure Prerequisites

GKE clusters must have multi-networking configured before deploying AICR bundles:

  • Multi-networking enabled (1 gVNIC + 4 RDMA NICs per a4x-highgpu-4g node)
  • Network + GKENetworkParamSet CRs for the gVNIC and 4 RDMA NICs (cluster-specific VPC/subnet values, but fixed object names; see below, not managed by AICR)
  • nccl-rdma-installer DaemonSet on GPU nodes (included in the AICR bundle)
  • Each GPUDirect-RDMA workload Pod must request all 4 GPUs and use all 4 RDMA NICs on a single node; RDMA can’t be shared between Pods on the same node (a GKE a4x-highgpu-4g constraint, not an AICR-specific one). AICR’s own recipes already request whole nodes this way; a custom workload built against this component must too.

The nccl-rdma-installer DaemonSet ships in the AICR bundle. The Network/ GKENetworkParamSet CRs and the multi-networking/VPC fabric underneath them are cluster provisioning: AICR’s gke-gb200-rdma health check detects them but does not create them.

Provisioning multi-networking

These steps are ordered, following Google’s A4X custom setup guide:

  1. Create the VPCs and subnets: two VPCs in the cluster’s region, one for the gVNIC (with one subnet) and one RDMA VPC (with four subnets, one per RDMA NIC); five subnets total across the two VPCs, not five separate VPCs.
  2. Create the cluster with multi-networking enabled (HIPPO’s GKECluster CR does this via spec.networks.managed.gb200NetworkStrategy).
  3. Apply the Network and GKENetworkParamSet CRs, one pair per NIC, binding each additional node network into the cluster so pods can reference it. Unlike TCPXO (see GKE TCPXO Networking), the object names are fixed, not cluster-specific: gvnic-1 for the gVNIC and rdma-0 through rdma-3 for the RDMA NICs. Only the vpc/vpcSubnet fields inside each GKENetworkParamSet vary per cluster (they name the VPC/subnet your cluster actually has):
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: gvnic-1
spec:
  vpc: "PREFIX-gvnic"
  vpcSubnet: "PREFIX-gvnic"
  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

Repeat for rdma-0 through rdma-3, pointing vpc at the single RDMA VPC from step 1 (the same value for all four) and vpcSubnet at that VPC’s four subnets (PREFIX-rdma-sub-0 through PREFIX-rdma-sub-3, or whatever names your subnets were given in step 1, with PREFIX replaced by your own), and set deviceMode: RDMA on all four, not NetDevice (that value is only correct for gvnic-1 above).

The fixed naming is a requirement, not a convention. AICR’s checks/gke-gb200-rdma/health-check.yaml asserts these five objects by exact name (gvnic-1, rdma-0..rdma-3), including spec.deviceMode and spec.parametersRef linkage. A cluster provisioned with different Network names passes Google’s own setup guide but fails this check; rename to match before running aicr validate.

  1. Create the GPU node pool on an a4x-highgpu-4g machine type, attaching the five network/subnet pairs as additionalNodeNetworkConfigs (the RDMA VPC repeated across its four subnets, plus the gVNIC VPC/subnet).

AICR installs the nccl-rdma-installer DaemonSet and detects the CRs; it does not provision the networking itself. These steps are a summary of the prerequisite AICR depends on, not a complete provisioning runbook; follow Google’s guide above for the full procedure, including firewall rules and supported GKE version floors.

Separately from GKE’s own networking version floor, all AICR GB200 GKE recipes (including gb200-gke-cos-training-slurm, which inherits it from gb200-gke-cos-training) enforce K8s.server.version >= 1.34: NVLS provisions the IMEX channel through a DRA ComputeDomain, which requires the GA resource.k8s.io/v1 API. aicr validate fails readiness on an older control plane with this constraint by name.

Verifying

kubectl get network.networking.gke.io \
  -o custom-columns='NAME:.metadata.name,PARAMETERS-REF:.spec.parametersRef.name'
kubectl get gkenetworkparamset.networking.gke.io \
  -o custom-columns='NAME:.metadata.name,DEVICE-MODE:.spec.deviceMode'

Expect gvnic-1 and rdma-0 through rdma-3 (the five prerequisite Networks from step 3), each bound to its GKENetworkParamSet via spec.parametersRef (shown in the PARAMETERS-REF column above). Fewer than five, or a GKENetworkParamSet with the wrong DEVICE-MODE, means the prerequisite is incomplete or misconfigured; aicr validate (via the gke-gb200-rdma health check) reports the shortfall by name.

You’ll also see a default network/GKENetworkParamSet pair in the same output; that one is GKE-managed (created automatically once multi-networking is enabled), not part of this prerequisite, and isn’t checked by name.

The shipped torch-distributed-rdma runtime

On the gb200-gke-cos-training-kubeflow recipe, AICR ships a pre-wired ClusterTrainingRuntime named torch-distributed-rdma (a sibling of torch-distributed, which stays as it is). It carries the networking.gke.io pod annotations mapping eth1 to gvnic-1 and eth2 through eth5 to rdma-0 through rdma-3, the /home/kubernetes/bin/nvidia and /home/kubernetes/bin/gib host mounts, NCCL_NET=gIB, IPC_LOCK, and the whole-node 4-GPU request, so a TrainJob references it by name and supplies no fabric configuration at all:

spec:
  runtimeRef:
    name: torch-distributed-rdma
    apiGroup: trainer.kubeflow.org
    kind: ClusterTrainingRuntime
  trainer:
    numNodes: 2                      # the shipped default; override freely
    image: my-registry/my-trainer:latest

Do not set resourcesPerNode. On the pinned Kubeflow Trainer v2.2.0 a TrainJob value replaces the runtime’s resource requirements outright rather than merging, so a job that sets it loses the nvidia.com/gpu: "4" request and with it the whole node. RDMA cannot be shared between Pods on a node, so that failure takes the fabric with it. If you must set it, repeat every resource in it, including the GPU request.

No recipe-generation input is required. gvnic-1 and rdma-0 through rdma-3 are fixed object names on every GB200/GKE cluster, so the runtime renders unconditionally and has no unset path that could ship it half-wired.

The runtime does not apply upstream’s NCCL tuning profile. The gIB installer lays that profile down on the host at /home/kubernetes/bin/gib/scripts/set_nccl_env.sh. A ClusterTrainingRuntime cannot source it, because Kubeflow Trainer owns the container command, so a workload that wants it sources the file from its own spec.trainer.command. AICR’s own NVLS performance validator does exactly that (see validators/performance/testdata/gb200/gke/runtime-nvls.yaml). Skipping it costs tuning, not correctness, because NCCL_NET=gIB from the runtime is what selects the fabric.

Overriding command replaces the torchrun invocation Trainer injects, so the override has to launch it. Pass no rendezvous flags. Trainer sets PET_NNODES, PET_NPROC_PER_NODE, PET_NODE_RANK, PET_MASTER_ADDR and PET_MASTER_PORT on the container, and torchrun reads each one as the default for the matching flag, so a hand-written flag can only contradict the job it was launched under.

spec:
  trainer:
    command: ["/bin/bash", "-c"]
    args:
      - |
        if [ -f /home/kubernetes/bin/gib/scripts/set_nccl_env.sh ]; then
          . /home/kubernetes/bin/gib/scripts/set_nccl_env.sh
        fi
        exec torchrun train.py

Driver Installer

a4x-highgpu-4g recipes generated with --profile gpuStack=bundle-installer (see GKE GPU Setup) get the driver from the bundle’s gcp-driver-installer component — Google’s cos-gpu-installer DaemonSet, deployed and versioned by AICR alongside the rest of the bundle, no manual DaemonSet apply required. This presumes the node-pool prerequisite (pools created with gpu-driver-version=disabled plus the gke-no-default-nvidia-gpu-device-plugin=true label) is already in place; the component’s nodeAffinity requires that label itself, so it never schedules onto a pool that hasn’t opted out of GKE’s managed install.

GB200-specific wrinkle: the component’s default partitionGpuImage (its partition-gpus init container, Google’s nvidia-partition-gpu MIG tool — a no-op here since this recipe allocates whole GPUs per node rather than configuring MIG) is pinned to an amd64-only digest and fails with exec format error on GB200’s arm64 nodes. AICR’s gb200-gke-cos-training and gb200-gke-cos-inference overlays (and everything based on them) set gcp-driver-installer.partitionGpuImage to a multi-arch digest automatically — found via live GB200 GKE validation, no action needed:

partitionGpuImage: "gcr.io/gke-release/nvidia-partition-gpu@sha256:de12f85ebfb4fb6c1893cd30c23aab662a72fa0448f97ef74fccb82d7522ef17"

The driver version itself (gcp-driver-installer.driverVersion) needs no GB200-specific override: the component’s default is COS-qualified for GB200 (see the component’s values.yaml for the qualified COS builds).

Validate the RDMA prerequisite before deploying

The gcp-driver-installer component deploys and orders itself ahead of the GPU Operator automatically within deploy.sh (see Driver Installer) — there’s nothing to hand-apply or pre-check for the driver anymore. The RDMA Network/GKENetworkParamSet CRs are still a true cluster-provisioning prerequisite applied before the node pool exists (see Provisioning multi-networking above); confirm they’re in place before running the bundle’s full deploy.sh with the kubectl commands in Verifying. The gke-gb200-rdma health check itself asserts the bundle’s nccl-rdma-installer DaemonSet is fully rolled out with healthy Pods, on top of the CRs, so it (and every other deployment-phase check) reports the bundle’s own components as missing until deploy.sh has actually deployed them.

Once deploy.sh has deployed the bundle, aicr validate --phase deployment against the bundle’s recipe.yaml (the component set it deployed) checks it, including gke-gb200-rdma, without running conformance or performance (see Validation):

aicr validate --recipe ./bundle/recipe.yaml --phase deployment

Storage Prerequisites

a4x-highgpu-4g nodes can’t attach Persistent Disk at all (regional or zonal, any type, including pd-balanced); only Hyperdisk. On a stock GKE Standard cluster the default StorageClass is standard-rwo (pd.csi.storage.gke.io, pd-balanced), but “default” isn’t inherent to GKE Standard itself: a cluster admin can repoint the storageclass.kubernetes.io/is-default-class annotation to any StorageClass. Run kubectl get storageclass first and check which one is annotated (default), its PROVISIONER, and (via kubectl get storageclass -o yaml) its parameters.type; don’t assume it’s standard-rwo/pd-balanced. Any PVC scheduled onto a GB200 node with no storageClassName set (which binds it to the cluster default) fails this way unless that default’s parameters.type is already Hyperdisk-backed: pd-balanced disk type cannot be used by a4x-highgpu-4g machine type (or the equivalent for whatever pd-* type the default actually provisions).

This includes the inference-perf validator’s model-weights cache PVC, which the populate Job pins to the a4x node. The gb200-gke-cos-inference-dynamo recipe provisions its own StorageClass for this, a4x-compatible (pd.csi.storage.gke.io, parameters.type: hyperdisk-balanced; see recipes/components/dynamo-platform/manifests/a4x-storage-class.yaml), and points the inference-model-cache-storage-class performance constraint at it. It is not the cluster default, so it has no effect on PVCs scheduled elsewhere, including CPU node pools that can’t attach Hyperdisk. No manual setup is needed for this PVC.

Override it only to reuse a differently named, already-Hyperdisk-backed StorageClass instead: set inference-model-cache-storage-class on the recipe overlay, or AICR_INFERENCE_PERF_MODEL_CACHE_STORAGE_CLASS=<name> on the inference-perf catalog entry’s env. AICR’s compatibility check accepts hyperdisk-balanced, hyperdisk-ml, or hyperdisk-extreme (see Validation).

When overriding to a StorageClass this bundle doesn’t own, for example one that already exists on the cluster under the name AICR would otherwise create, also pass --set dynamo-platform:a4xStorageClass.create=false at bundle time, so AICR doesn’t render the a4x-compatible manifest at all.

Running the NCCL Benchmark

The GB200 GKE training recipe (gb200-gke-cos-training) selects the NVLS-variant performance check (nccl-all-reduce-bw-nvls): MNNVL across the A4X nodes’ IMEX domain is the fabric that carries all-reduce traffic; gIB is the transport driver underneath, not the NCCL algorithm itself. Run it via:

aicr validate --recipe recipes/overlays/gb200-gke-cos-training.yaml \
  --phase performance

References