A3 Edge を使用する GKE クラスタで NCCL を実行する

このページでは、GPUDirect-TCPX ネットワーキング プロトコルで A3 Edge を使用するカスタム GKE クラスタで NVIDIA Collective Communications Library(NCCL)テストを実行する方法について説明します。カスタム GKE クラスタは、gcloud コマンドを使用して作成するクラスタです。

このページで説明するテストは、次のシナリオで使用できます。

始める前に

このページのテストでは、JobSet と トポロジを考慮したスケジューリング(TAS)を使用した Kueue を使用します。テストを実行する前に、クラスタを設定して次の操作を行う必要があります。

  1. JobSet をインストールします。

  2. Kueue をインストールします。

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/kueue/releases/download/v0.16.5/manifests.yaml
    

Jobset と Kueue を使用してクラスタを設定する

JobSet と Kueue をインストールしたら、次の手順を行います。

  1. 次のマニフェストを kueue-config.yaml ファイルとして保存します。

    apiVersion: kueue.x-k8s.io/v1beta2
    kind: Topology
    metadata:
      name: "gke-default"
    spec:
      levels:
      - nodeLabel: "cloud.google.com/gce-topology-block"
      - nodeLabel: "cloud.google.com/gce-topology-subblock"
      - nodeLabel: "cloud.google.com/gce-topology-host"
      - nodeLabel: "kubernetes.io/hostname"
    ---
    apiVersion: kueue.x-k8s.io/v1beta2
    kind: ResourceFlavor
    metadata:
      name: a3-edge-flavor
    spec:
      nodeLabels:
        cloud.google.com/gke-accelerator: nvidia-h100-80gb
      topologyName: "gke-default"
    ---
    apiVersion: kueue.x-k8s.io/v1beta2
    kind: ResourceFlavor
    metadata:
      name: a3-edge-dws-flavor
    spec:
      nodeLabels:
        cloud.google.com/gke-accelerator: nvidia-h100-80gb
      topologyName: "gke-default"
      tolerations:
      - key: "cloud.google.com/gke-queued"
        operator: "Exists"
        effect: NoSchedule
    ---
    apiVersion: kueue.x-k8s.io/v1beta2
    kind: AdmissionCheck
    metadata:
      name: dws-prov
    spec:
      controllerName: kueue.x-k8s.io/provisioning-request
      parameters:
        apiGroup: kueue.x-k8s.io
        kind: ProvisioningRequestConfig
        name: dws-config
    ---
    apiVersion: kueue.x-k8s.io/v1beta2
    kind: ProvisioningRequestConfig
    metadata:
      name: dws-config
    spec:
      provisioningClassName: queued-provisioning.gke.io
      podSetUpdates:
      - key: autoscaling.gke.io/provisioning-request
      valueFromProvisioningClassDetail: ResizeRequestName
      managedResources:
      - nvidia.com/gpu
    ---
    apiVersion: kueue.x-k8s.io/v1beta2
    kind: ClusterQueue
    metadata:
      name: cq-tas
    spec:
      namespaceSelector: {}
      clusterQueueingStrategy: BestEffortFIFO
      resourceGroups:
      - flavors:
        - name: a3-edge-flavor
          resources:
          - name: "cpu"
            nominalQuota: 1000
          - name: "memory"
            nominalQuota: 1000Ti
          - name: "nvidia.com/gpu"
            nominalQuota: 1000
        - name: a3-edge-dws-flavor
          resources:
          - name: "cpu"
            nominalQuota: 1000
          - name: "memory"
            nominalQuota: 1000Ti
          - name: "nvidia.com/gpu"
            nominalQuota: 1000
      admissionChecksStrategy:
        admissionChecks:
        - name: "dws-prov"
          onFlavors: [a3-edge-dws-flavor]
    ---
    apiVersion: kueue.x-k8s.io/v1beta2
    kind: LocalQueue
    metadata:
      namespace: default
      name: lq-tas
    spec:
      clusterQueue: cq-tas
    

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

    kubectl apply -f kueue-config.yaml
    

トポロジ対応スケジューリング(TAS)が有効になっているワークロードを実行する場合は、ワークロード マニフェストで次のいずれかのアノテーションを使用して、トポロジ制約の適用を厳密に指定できます。

  • kueue.x-k8s.io/podset-required-topology: このアノテーションを使用すると、Kueue は、リクエストされたトポロジ制約内でワークロードをスケジュールできるようになるまで、スケジューリングをブロックします。このアノテーションを使用すると、最適なパフォーマンスを得るために Pod が一緒に配置されます。

  • kueue.x-k8s.io/podset-preferred-topology: このアノテーションを使用すると、Kueue はリクエストされたトポロジ制約内で Pod のスケジュールを試みますが、それが不可能な場合は、トポロジ制約を満たさずにワークロードを承認します。

どちらのアノテーションでも、トポロジ制約として次のいずれかの値を指定します。

  • cloud.google.com/gce-topology-block: 同じネットワーク ブロック内の Pod をスケジュールします。
  • cloud.google.com/gce-topology-subblock: 同じラック内の Pod をスケジュールします。
  • cloud.google.com/gce-topology-host: 同じ物理ホストに Pod をスケジュールします。

2 つの Flex Start ノードでテストする

A3 Edge Flex Start VM を使用する GKE クラスタで NCCL テストを実行するには、次の手順を使用します。この手順では、JobSet マニフェストを使用して、2 つのノードで NCCL テストを実行します。

  1. 次のマニフェストを nccl-tas-jobset.yaml ファイルとして保存します。

    apiVersion: v1
    kind: ConfigMap
    metadata:
      name: nccl-config
    data:
      allgather.sh: |
        #!/bin/bash
        for script in /configs/*; do
          name=$(basename $script)
          cp $script "/scripts/$name"
          chmod +x "/scripts/$name"
        done
        /scripts/init_ssh.sh ${@};
        pushd /scripts;
        /scripts/gen_hostfiles.sh ${@};
        popd;
        /scripts/run-allgather.sh 8 eth1,eth2,eth3,eth4 1M 512M ${#};
    ---
    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: nccl-tas-test
      labels:
        kueue.x-k8s.io/queue-name: lq-tas
    spec:
      suspend: true
      network:
        enableDNSHostnames: true
      replicatedJobs:
      - name: worker
        replicas: 2
        template:
          spec:
            parallelism: 1
            completions: 1
            template:
              metadata:
                annotations:
                  kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-block"
                  networking.gke.io/default-interface: 'eth0'
                  networking.gke.io/interfaces: |
                            [
                              {"interfaceName":"eth0","network":"default"},
                              {"interfaceName":"eth1","network":"vpc0"},
                              {"interfaceName":"eth2","network":"vpc1"},
                              {"interfaceName":"eth3","network":"vpc2"},
                              {"interfaceName":"eth4","network":"vpc3"}
                            ]
              spec:
                terminationGracePeriodSeconds: 0
                nodeSelector:
                  cloud.google.com/gke-accelerator: nvidia-h100-80gb
                tolerations:
                - key: cloud.google.com/gke-queued
                  effect: NoSchedule
                  value: "true"
                - key: "nvidia.com/gpu"
                  operator: "Exists"
                  effect: "NoSchedule"
                setHostnameAsFQDN: true
                containers:
                - name: tcpx-daemon
                  image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/tcpgpudmarxd-dev:v2.0.11
                  command:
                    - /tcpgpudmarxd/build/app/tcpgpudmarxd
                    - --gpu_nic_preset
                    - a3vm
                    - --gpu_shmem_type
                    - fd
                    - --uds_path
                    - /run/tcpx
                    - --setup_param
                    - "--verbose 128 2 0 "
                  securityContext:
                    privileged: true
                    capabilities:
                      add:
                        - NET_ADMIN
                  volumeMounts:
                    - name: libraries
                      mountPath: /usr/local/nvidia/lib64
                    - name: tcpx-socket
                      mountPath: /run/tcpx
                    - name: sys
                      mountPath: /hostsysfs
                    - name: proc-sys
                      mountPath: /hostprocsysfs
                  env:
                    - name: LD_LIBRARY_PATH
                      value: /usr/local/nvidia/lib64
                - name: nccl-test
                  image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx-dev:v3.1.8
                  command:
                    - bash
                    - -c
                    - |
                      /scripts/container_entry.sh daemon;
                      sleep infinity;
                  securityContext:
                    privileged: true
                  volumeMounts:
                    - name: tcpx-socket
                      mountPath: /tmp
                    - name: libraries
                      mountPath: /usr/local/nvidia/lib64
                    - name: nccl-config
                      mountPath: /configs
                    - name: shared-memory
                      mountPath: /dev/shm
                  resources:
                    limits:
                      cpu: "200"
                      memory: "1800Gi"
                      nvidia.com/gpu: 8
                    requests:
                      cpu: "200"
                      memory: "1800Gi"
                      nvidia.com/gpu: 8
                volumes:
                - name: libraries
                  hostPath:
                    path: /home/kubernetes/bin/nvidia/lib64
                - name: tcpx-socket
                  emptyDir: {}
                - name: sys
                  hostPath:
                    path: /sys
                - name: proc-sys
                  hostPath:
                    path: /proc/sys
                - name: shared-memory
                  emptyDir:
                    medium: Memory
                    sizeLimit: 250Gi
                - name: nccl-config
                  configMap:
                    name: nccl-config
                    defaultMode: 0777
    

  2. マニフェストをクラスタに適用します。

    kubectl apply -f nccl-tas-jobset.yaml
    
  3. JobSet が許可され、実行されていることを確認します。

    kubectl get jobset nccl-tas-test
    

    JobSet の一時停止が解除され、Pod が Running ステータスになるまで待ちます。

  4. 最初のワーカー Pod から allgather.sh スクリプトを実行して、NCCL テストをトリガーします。

    kubectl exec --stdin --tty --container=nccl-test nccl-tas-test-worker-0-0 -- /configs/allgather.sh nccl-tas-test-worker-0-0 nccl-tas-test-worker-1-0
    

    2 ノードのテストの出力は次のようになります。

      #                                                              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)
          1048576         16384     float    none      -1    696.8    1.50    1.41      0    729.0    1.44    1.35      0
          ...
          536870912       8388608     float    none      -1   7101.7   75.60   70.87      0   7060.9   76.03   71.28      0
      # Out of bounds values : 0 OK
      # Avg bus bandwidth    : 29.8293
    

TAS を使用して NCCL テスト ワークロードをデプロイする

ノードが 3 つ以上ある場合は、(TAS) を使用する次のテストを使用することをおすすめします。A3 Edge Flex Start VM を使用する GKE クラスタで TAS を使用して NCCL テストを実行するには、次の手順を行います。

  1. 次のマニフェストを nccl-jobset-test.yaml ファイルとして保存します。NUM_NODES は、ノードプール内のノード数に置き換えます。

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: nccl-ag
      labels:
        kueue.x-k8s.io/queue-name: lq-tas
    spec:
      ttlSecondsAfterFinished: 1200
      suspend: true
      network:
        enableDNSHostnames: true
      replicatedJobs:
        - name: worker
          template:
            spec:
              parallelism: NUM_NODES
              completions: NUM_NODES
              template:
                metadata:
                  annotations:
                    kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-subblock"
                    networking.gke.io/default-interface: 'eth0'
                    networking.gke.io/interfaces: |
                            [
                              {"interfaceName":"eth0","network":"default"},
                              {"interfaceName":"eth1","network":"vpc0"},
                              {"interfaceName":"eth2","network":"vpc1"},
                              {"interfaceName":"eth3","network":"vpc2"},
                              {"interfaceName":"eth4","network":"vpc3"}
                            ]
                spec:
                  activeDeadlineSeconds: 3600
                  restartPolicy: Never
                  nodeSelector:
                    cloud.google.com/gke-accelerator: nvidia-h100-80gb
                  tolerations:
                  - key: cloud.google.com/gke-queued
                    operator: "Exists"
                    effect: NoSchedule
                  - key: "nvidia.com/gpu"
                    operator: "Exists"
                    effect: "NoSchedule"
                  setHostnameAsFQDN: true
                  volumes:
                  - name: proc
                    hostPath:
                      path: /proc
                  - name: nvidia
                    hostPath:
                      path: /home/kubernetes/bin/nvidia
                  - name: libraries
                    hostPath:
                      path: /home/kubernetes/bin/nvidia/lib64
                  - name: tcpx-socket
                    emptyDir: {}
                  - name: shared-memory
                    emptyDir:
                      medium: "Memory"
                      sizeLimit: 250Gi
                  containers:
                  - name: tcpx-daemon
                    image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/tcpgpudmarxd-dev:v2.0.11
                    command:
                      - /tcpgpudmarxd/build/app/tcpgpudmarxd
                      - --gpu_nic_preset
                      - a3vm
                      - --gpu_shmem_type
                      - fd
                      - --uds_path
                      - /run/tcpx
                      - --setup_param
                      - "--verbose 128 2 0 "
                    securityContext:
                      privileged: true
                    volumeMounts:
                      - name: tcpx-socket
                        mountPath: /run/tcpx
                      - name: libraries
                        mountPath: /usr/local/nvidia/lib64
                  - name: nccl-test
                    stdin: true
                    tty: true
                    image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx-dev:v3.1.8
                    securityContext:
                      privileged: true
                    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
                    - |
                      /scripts/container_entry.sh daemon &
                      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 i in `seq 0 $(($N_NODES-1))`; do
                        OTHER=${WORKERS_BASENAME}-${i}.${POSTFIX}
                        until ssh -p 222 -o StrictHostKeyChecking=no $OTHER hostname; do
                          sleep 10
                        done
                        echo ${OTHER} port=222 slots=8 | tee -a /tmp/hostfile;
                      done
                      if [[ "${NODE_RANK}" -eq "0" ]]; then
                          /scripts/run-allgather.sh 8 eth1,eth2,eth3,eth4 1M 512M ${N_NODES}
                      else
                          while ping -c 1 ${WORKERS_BASENAME}-0.${POSTFIX}; do
                          sleep 5
                      done
                      fi
                      exit 0
                    volumeMounts:
                    - name: nvidia
                      mountPath: /usr/local/nvidia
                    - name: tcpx-socket
                      mountPath: /tmp
                    - name: libraries
                      mountPath: /usr/local/nvidia/lib64
                    - name: shared-memory
                      mountPath: /dev/shm
                    resources:
                      limits:
                        cpu: "200"
                        memory: "1800Gi"
                        nvidia.com/gpu: 8
                      requests:
                        cpu: "200"
                        memory: "1800Gi"
                        nvidia.com/gpu: 8
    

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

    kubectl apply -f nccl-jobset-test.yaml
    
  3. ワークロードが許可され、Completed 状態になっていることを確認します。

  4. nccl-ag-worker-0-0-.* に一致する Pod のログを取得して、結果を確認します。

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

次のステップ