Run NCCL on custom GKE clusters that use A4X

This page describes how to run NCCL/gIB tests on provisioned clusters that use GPUDirect RDMA. It describes tests for the following scenarios:

Test on two nodes

  1. Connect to your cluster:

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

    Replace the following variables:

    • CLUSTER_NAME: the name of your cluster, which, for the clusters created with Cluster Toolkit, is based on the DEPLOYMENT_NAME.
    • COMPUTE_REGION: the name of the compute region.
  2. To deploy an NCCL test workload of two test Pods that are running on two A4X nodes, run the following:

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/refs/heads/master/gpudirect-rdma/nccl-test-imex-a4x.yaml
    
  3. Check if the Pods are both running on some nodes:

    kubectl get pods nccl-test-host-1 nccl-test-host-2
    

    If the two Pods show a Running status, you can proceed to the next step.

  4. Trigger an all-gather test for the A4X nodes:

    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
    

    The output is similar to the following:

    #                                                              out-of-place                       in-place
    #       size         count      type   redop    root     time   algbw   busbw #wrong     time   algbw   busbw #wrong
    #        (B)    (elements)                               (us)  (GB/s)  (GB/s)            (us)  (GB/s)  (GB/s)
            1024            32     float    none      -1    21.20    0.05    0.04      0    20.56    0.05    0.04      0
            2048            64     float    none      -1    21.03    0.10    0.09      0    20.82    0.10    0.09      0
            4096           128     float    none      -1    21.11    0.19    0.17      0    20.98    0.20    0.17      0
            8192           256     float    none      -1    21.51    0.38    0.33      0    21.15    0.39    0.34      0
           16384           512     float    none      -1    21.85    0.75    0.66      0    21.72    0.75    0.66      0
           32768          1024     float    none      -1    24.08    1.36    1.19      0    23.73    1.38    1.21      0
           65536          2048     float    none      -1    24.68    2.66    2.32      0    24.02    2.73    2.39      0
          131072          4096     float    none      -1    24.93    5.26    4.60      0    24.30    5.40    4.72      0
          262144          8192     float    none      -1    24.86   10.55    9.23      0    24.33   10.78    9.43      0
          524288         16384     float    none      -1    25.10   20.89   18.28      0    24.48   21.41   18.74      0
         1048576         32768     float    none      -1    25.43   41.24   36.09      0    24.82   42.25   36.97      0
         2097152         65536     float    none      -1    32.30   64.93   56.81      0    31.28   67.04   58.66      0
         4194304        131072     float    none      -1    45.92   91.34   79.92      0    44.22   94.84   82.99      0
         8388608        262144     float    none      -1    71.38  117.52  102.83      0    68.98  121.61  106.41      0
        16777216        524288     float    none      -1    74.17  226.20  197.93      0    72.37  231.83  202.85      0
        33554432       1048576     float    none      -1    116.6  287.84  251.86      0    112.7  297.75  260.54      0
        67108864       2097152     float    none      -1    188.9  355.27  310.86      0    184.0  364.71  319.12      0
       134217728       4194304     float    none      -1    309.6  433.56  379.36      0    299.7  447.83  391.85      0
       268435456       8388608     float    none      -1    559.0  480.23  420.20      0    540.3  496.85  434.75      0
       536870912      16777216     float    none      -1   1053.7  509.52  445.83      0   1021.4  525.64  459.93      0
      1073741824      33554432     float    none      -1   2087.4  514.39  450.10      0   2013.8  533.19  466.54      0
      2147483648      67108864     float    none      -1   4154.7  516.88  452.27      0   3987.4  538.57  471.25      0
      4294967296     134217728     float    none      -1   8289.2  518.14  453.37      0   7907.4  543.16  475.26      0
      8589934592     268435456     float    none      -1    16556  518.85  453.99      0    15726  546.24  477.96      0
    # Out of bounds values : 0 OK
    # Avg bus bandwidth    : 175.233
    #
    

Test with TAS

To validate the functionality of the provisioned cluster, you can run the following NCCL test with TAS.

Configure Kueue with TAS enabled

  1. Install Kueue with TAS enabled.
  2. Configure Kueue with TAS enabled by creating the following file, which you name a4x-kueue-config.yaml:

    apiVersion: kueue.x-k8s.io/v1alpha1
    kind: Topology
    metadata:
      name: "a4x-default"
    spec:
      levels:
      - nodeLabel: "cloud.google.com/gce-topology-block"
      - nodeLabel: "cloud.google.com/gce-topology-subblock"
      - nodeLabel: "cloud.google.com/gke-nodepool"
      - nodeLabel: "cloud.google.com/gce-topology-host"
      - nodeLabel: "kubernetes.io/hostname"
    ---
    kind: ResourceFlavor
    apiVersion: kueue.x-k8s.io/v1beta1
    metadata:
      name: "a4x"
    spec:
      nodeLabels:
        cloud.google.com/gke-accelerator: nvidia-gb200
      topologyName: "a4x-default"
      tolerations:
      - key: "nvidia.com/gpu"
        operator: "Exists"
        effect: NoSchedule
      - key: "kubernetes.io/arch"
        operator: "Exists"
        effect: NoSchedule
    ---
    apiVersion: kueue.x-k8s.io/v1beta1
    kind: ClusterQueue
    metadata:
      name: "a4x"
    spec:
      namespaceSelector: {} # match all.
      resourceGroups:
      - coveredResources: ["nvidia.com/gpu"]
        flavors:
        - name: "a4x"
          resources:
          - name: "nvidia.com/gpu"
            nominalQuota: 1_000_000_000
    ---
    apiVersion: kueue.x-k8s.io/v1beta1
    kind: LocalQueue
    metadata:
      namespace: "default"
      name: "a4x"
    spec:
      clusterQueue: "a4x"
    
  3. Run the test:

    kubectl apply -f a4x-kueue-config.yaml
    

Schedule a topology-aware NCCL test with Kueue with TAS enabled

The following workload must be placed within a single NVLink Domain sub-block.

  1. Install JobSet, a Kubernetes-native API for managing of group of Kubernetes Jobs as a unit. Ensure that your non-GPU node pools have enough resources to schedule the JobSet controllers.
  2. Create the following file with the name nccl-tas-test.yaml. Replace NUM_NODES with the intended number of nodes to run the NCCL test, up to 18:

    apiVersion: resource.nvidia.com/v1beta1
    kind: ComputeDomain
    metadata:
      name: nccl-test-compute-domain
    spec:
      numNodes: NUM_NODES
      channel:
        resourceClaimTemplate:
          name: nccl-test-compute-domain-channel
    ---
    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: kueue-tas-nccl-all-gather
      labels:
        kueue.x-k8s.io/queue-name: a4x
    spec:
      ttlSecondsAfterFinished: 1200
      network:
        enableDNSHostnames: true
      replicatedJobs:
        - name: worker
          template:
            spec:
              parallelism: NUM_NODES
              completions: NUM_NODES
              template:
                metadata:
                  annotations:
                    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-subblock"
                    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:
                  activeDeadlineSeconds: 3600
                  restartPolicy: Never
                  nodeSelector:
                    cloud.google.com/gke-accelerator: nvidia-gb200
                  tolerations:
                  - key: nvidia.com/gpu
                    operator: Equal
                    value: present
                    effect: NoSchedule
                  - key: kubernetes.io/arch
                    operator: Equal
                    value: arm64
                    effect: NoSchedule
                  setHostnameAsFQDN: true
                  volumes:
                  - name: gib
                    hostPath:
                      path: /home/kubernetes/bin/gib
                  - name: nvidia
                    hostPath:
                      path: /home/kubernetes/bin/nvidia
                  - name: lib64
                    hostPath:
                      path: /lib64
                  - name: shared-memory
                    emptyDir:
                      medium: "Memory"
                      sizeLimit: 250Gi
                  resourceClaims:
                  - name: compute-domain-channel
                    resourceClaimTemplateName: nccl-test-compute-domain-channel
                  containers:
                  - name: nccl-test
                    stdin: true
                    tty: true
                    image: us-docker.pkg.dev/gce-ai-infra/gpudirect-gib/nccl-plugin-gib-diagnostic-arm64:v1.0.4
                    env:
                    - name: MY_NODE_NAME
                      valueFrom:
                        fieldRef:
                          fieldPath: spec.nodeName
                    - name: OMPI_ALLOW_RUN_AS_ROOT
                      value: "1"
                    - name: OMPI_ALLOW_RUN_AS_ROOT_CONFIRM
                      value: "1"
                    - name: N_NODES
                      value: "NUM_NODES"
                    - name: LD_LIBRARY_PATH
                      value: /usr/local/nvidia/lib64
                    command:
                    - bash
                    - -c
                    - |
                      set -x
                      echo "Starting workload container on ${MY_NODE_NAME} for $N_NODES benchmark"
                      # Install ping
                      apt update -y
                      apt install -y iputils-ping
    
                      # Start sshd
                      /scripts/container_entry.sh daemon &
    
                      # Get helper variables to form all hostnames
                      export POSTFIX=$(hostname | cut -d . -f 2-)
                      export WORKERS_BASENAME=$(hostname | cut -d . -f 1 | rev | cut -d - -f 2- | rev )
                      export NODE_RANK=$JOB_COMPLETION_INDEX
    
                      # For every worker, wait till online and add to hostfile
                      for i in `seq 0 $(($N_NODES-1))`; do
                        OTHER=${WORKERS_BASENAME}-${i}.${POSTFIX}
                        until ssh -p 222 -o StrictHostKeyChecking=no $OTHER hostname; do
                          echo Waiting for ${OTHER}...
                          sleep 10
                        done
                        echo ${OTHER} port=222 slots=4 | tee -a /tmp/hostfile;
                      done
    
                      cat /tmp/hostfile
    
                      # Launch from head node
                      if [[ "${NODE_RANK}" -eq "0" ]]; then
    
                          # World Level = 0x0, Rail Aligned = 0x7
                          export NCCL_TESTS_SPLIT_MASK="0x0";
    
                          # Force use of libnccl-gib
                          export NCCL_NET=gIB
    
                          # Set all the correct libnccl-gib environment variables
                          source /usr/local/gib/scripts/set_nccl_env.sh
    
                          # Get all relevant NCCL / env vars to pass to all workers
                          ENV_VARS=$(echo ${!NCCL*} ${!OMPI*} LD_LIBRARY_PATH PATH | sed 's/ / -x /g')
    
                          mpirun --hostfile /tmp/hostfile \
                            -x $ENV_VARS  \
                            -mca plm_rsh_no_tree_spawn 1 \
                            --mca orte_keep_fqdn_hostnames 1 \
                            --mca btl self,tcp \
                            --mca btl_tcp_if_include eth0 \
                            --bind-to none \
                            --mca plm_rsh_agent "ssh -q -o LogLevel=ERROR -o StrictHostKeyChecking=no -p 222" \
                            /third_party/nccl-tests/build/all_gather_perf -b 1K -e 8G -f 2 -g 1 -w 5 --iters 100 -c 1
    
                      else
                          while ping -c 1 ${WORKERS_BASENAME}-0.${POSTFIX}; do
                          sleep 5
                      done
                      fi
    
                      exit 0
                    volumeMounts:
                    - name: nvidia
                      mountPath: /usr/local/nvidia
                    - name: gib
                      mountPath: /usr/local/gib
                    - name: shared-memory
                      mountPath: /dev/shm
                    resources:
                      limits:
                        nvidia.com/gpu: 4
                      requests:
                        nvidia.com/gpu: 4
                      claims:
                        - name: compute-domain-channel
                  restartPolicy: Never
    
  3. Run the test:

    kubectl apply -f nccl-tas-test.yaml
    
  4. Check the test result by reviewing the logs:

    kubectl logs $(kubectl get pods -o go-template='{{range .items}}{{.metadata.name}}{{"\n"}}{{end}}' | grep kueue-tas-nccl-all-gather-worker-0-0)

    The output should be similar to the following:

     #       size         count      type   redop    root     time   algbw   busbw #wrong     time   algbw   busbw #wrong
     #        (B)    (elements)                               (us)  (GB/s)  (GB/s)            (us)  (GB/s)  (GB/s)
             1024             8     float    none      -1    56.72    0.02    0.02      0    56.12    0.02    0.02      0
             2048            16     float    none      -1    56.85    0.04    0.03      0    56.87    0.04    0.03      0
             4096            32     float    none      -1    57.53    0.07    0.07      0    57.47    0.07    0.07      0
             8192            64     float    none      -1    58.43    0.14    0.14      0    58.27    0.14    0.14      0
            16384           128     float    none      -1    59.29    0.28    0.27      0    58.87    0.28    0.27      0
            32768           256     float    none      -1    60.02    0.55    0.53      0    59.60    0.55    0.53      0
            65536           512     float    none      -1    61.83    1.06    1.03      0    61.64    1.06    1.03      0
           131072          1024     float    none      -1    70.99    1.85    1.79      0    70.82    1.85    1.79      0
           262144          2048     float    none      -1    71.56    3.66    3.55      0    71.07    3.69    3.57      0
           524288          4096     float    none      -1    72.62    7.22    6.99      0    71.90    7.29    7.06      0
          1048576          8192     float    none      -1    72.80   14.40   13.95      0    72.31   14.50   14.05      0
          2097152         16384     float    none      -1    73.40   28.57   27.68      0    72.96   28.74   27.85      0
          4194304         32768     float    none      -1    73.86   56.78   55.01      0    73.44   57.12   55.33      0
          8388608         65536     float    none      -1    102.5   81.86   79.30      0    101.4   82.69   80.11      0
         16777216        131072     float    none      -1    158.3  105.97  102.66      0    156.8  107.02  103.68      0
         33554432        262144     float    none      -1    158.4  211.89  205.26      0    157.5  212.99  206.33      0
         67108864        524288     float    none      -1    250.7  267.68  259.32      0    248.7  269.81  261.38      0
        134217728       1048576     float    none      -1    417.7  321.29  311.25      0    414.1  324.13  314.01      0
        268435456       2097152     float    none      -1    728.8  368.32  356.81      0    721.5  372.08  360.45      0
        536870912       4194304     float    none      -1   1226.5  437.72  424.04      0   1216.1  441.46  427.66      0
       1073741824       8388608     float    none      -1   2268.4  473.35  458.56      0   2247.0  477.86  462.93      0
       2147483648      16777216     float    none      -1   4330.6  495.88  480.39      0   4291.6  500.39  484.76      0
       4294967296      33554432     float    none      -1   8640.9  497.05  481.52      0   8544.0  502.69  486.98      0
       8589934592      67108864     float    none      -1    17258  497.75  482.19      0    17052  503.75  488.00      0
     # Out of bounds values : 0 OK
     # Avg bus bandwidth    : 157.091
    

What's next