Executar o NCCL em clusters personalizados do GKE que usam o A4X Max

Este documento descreve como executar testes de NCCL/gIB em clusters provisionados que usam GPUDirect RDMA. Dependendo do seu caso de uso, use uma das seguintes opções:

Teste em dois nós

O teste a seguir executa uma carga de trabalho do NCCL em dois nós. Entenda o seguinte sobre este teste:

  • Por padrão, o GKE programa os dois pods para pools de nós separados, se disponíveis. Se os pools de nós forem criados com domínios NVLink distintos, esse teste representará a capacidade de processamento de RDMA entre domínios. Para programar pods no mesmo domínio, modifique a afinidade de pods para programar no mesmo pool de nós.

Execute o teste de dois nós:

  1. Conecte-se ao cluster:

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

    Substitua as seguintes variáveis:

    • CLUSTER_NAME: o nome do seu cluster, que, para os clusters criados com o Cluster Toolkit, é baseado no DEPLOYMENT_NAME.
    • COMPUTE_REGION: o nome da região de computação.
  2. Implante uma carga de trabalho de teste da NCCL de dois pods de teste em execução em dois nós A4X Max:

    kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-test-a4x-max.yaml
    
  3. Verifique se os pods estão sendo executados em alguns nós:

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

    Se os dois pods mostrarem o status Running, siga para a próxima etapa.

  4. Acione um teste de todos os nós A4X Max:

    HOSTS="nccl-host-1 nccl-host-2"
    kubectl exec nccl-test-host-1 -it -- bash -c "/usr/local/gib/scripts/run_nccl_tests.sh -t alltoall -b 1M -e 16G ${HOSTS}"
    

    O resultado será o seguinte:

    # nccl-tests version 2.17.6 nccl-headers=22809 nccl-library=22809
    # Collective test starting: all_gather_perf
    # nThread 1 nGpus 1 minBytes 4096 maxBytes 17179869184 step: 2(factor) warmup iters: 1 iters: 20 agg iters: 1 validation: 1 graph: 0
    #
    # Using devices
    #  Rank  0 Group  0 Pid    299 on gke-psch-ug-cluster-a4xmax-2-1391c2ef-m9bm device  0 [0008:06:00] NVIDIA GB300
    #  Rank  1 Group  0 Pid    300 on gke-psch-ug-cluster-a4xmax-2-1391c2ef-m9bm device  1 [0009:06:00] NVIDIA GB300
    #  Rank  2 Group  0 Pid    301 on gke-psch-ug-cluster-a4xmax-2-1391c2ef-m9bm device  2 [0018:06:00] NVIDIA GB300
    #  Rank  3 Group  0 Pid    302 on gke-psch-ug-cluster-a4xmax-2-1391c2ef-m9bm device  3 [0019:06:00] NVIDIA GB300
    #  Rank  4 Group  0 Pid    237 on gke-psch-ug-cluster-a4xmax-03878a77-d3lp device  0 [0008:06:00] NVIDIA GB300
    #  Rank  5 Group  0 Pid    238 on gke-psch-ug-cluster-a4xmax-03878a77-d3lp device  1 [0009:06:00] NVIDIA GB300
    #  Rank  6 Group  0 Pid    239 on gke-psch-ug-cluster-a4xmax-03878a77-d3lp device  2 [0018:06:00] NVIDIA GB300
    #  Rank  7 Group  0 Pid    240 on gke-psch-ug-cluster-a4xmax-03878a77-d3lp device  3 [0019:06:00] NVIDIA GB300
    #
    #                                                              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)
            4096           128     float    none      -1    28.51    0.14    0.13       0    27.71    0.15    0.13       0
            8192           256     float    none      -1    28.10    0.29    0.26       0    28.40    0.29    0.25       0
           16384           512     float    none      -1    28.55    0.57    0.50       0    28.19    0.58    0.51       0
           32768          1024     float    none      -1    30.56    1.07    0.94       0    29.65    1.11    0.97       0
           65536          2048     float    none      -1    33.30    1.97    1.72       0    33.14    1.98    1.73       0
          131072          4096     float    none      -1    36.18    3.62    3.17       0    36.14    3.63    3.17       0
          262144          8192     float    none      -1    38.50    6.81    5.96       0    94.91    2.76    2.42       0
          524288         16384     float    none      -1   152.25    3.44    3.01       0    54.79    9.57    8.37       0
         1048576         32768     float    none      -1    63.82   16.43   14.38       0    64.06   16.37   14.32       0
         2097152         65536     float    none      -1    65.10   32.21   28.19       0    66.13   31.71   27.75       0
         4194304        131072     float    none      -1    67.73   61.92   54.18       0    67.16   62.45   54.65       0
         8388608        262144     float    none      -1    79.65  105.31   92.15       0    80.02  104.83   91.73       0
        16777216        524288     float    none      -1   189.74   88.42   77.37       0   187.57   89.44   78.26       0
        33554432       1048576     float    none      -1   252.85  132.70  116.11       0   202.31  165.86  145.13       0
        67108864       2097152     float    none      -1   250.55  267.85  234.37       0   276.11  243.06  212.67       0
       134217728       4194304     float    none      -1   394.38  340.33  297.79       0   487.60  275.26  240.85       0
       268435456       8388608     float    none      -1   717.97  373.88  327.15       0   799.98  335.55  293.61       0
       536870912      16777216     float    none      -1  1421.29  377.73  330.52       0  1392.81  385.46  337.28       0
      1073741824      33554432     float    none      -1  2783.37  385.77  337.55       0  2596.97  413.46  361.78       0
      2147483648      67108864     float    none      -1  5396.10  397.97  348.22       0  5059.01  424.49  371.43       0
      4294967296     134217728     float    none      -1  10579.7  405.96  355.22       0  9918.44  433.03  378.90       0
      8589934592     268435456     float    none      -1  21012.9  408.79  357.69       0  20043.4  428.57  375.00       0
     17179869184     536870912     float    none      -1  42091.7  408.15  357.13       0  40243.2  426.90  373.54       0
    # Out of bounds values : 0 OK
    # Avg bus bandwidth    : 146.047
    #
    # Collective test concluded: all_gather_perf
    

    Se os pods forem programados em nós em domínios NVLink distintos, esse teste representará a taxa de transferência de RDMA entre domínios, conforme mostrado na saída fornecida. Para distribuir entre pools de nós criados em domínios NVLink distintos, modifique a afinidade de especificação do pod em nccl-test-a4x-max.yaml com o seguinte:

    spec:
    ...
    affinity:
    podAntiAffinity:
      preferredDuringSchedulingIgnoredDuringExecution:
      - weight: 100
        podAffinityTerm:
          labelSelector:
            matchLabels:
              app: nccl-test
          topologyKey: cloud.google.com/gke-nodepool
    

Testar usando JobSet

  1. Instalar JobSet:

    VERSION=v0.10.1
    kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/$VERSION/manifests.yaml
    
  2. Verifique se os pools de nós sem GPU têm recursos suficientes para programar os controladores JobSet. Siga a etapa para definir seus próprios ajustes de recursos.

    Para mais informações sobre a instalação do JobSet, consulte Instalação.

  3. Execute os comandos a seguir, substituindo NUM_NODES pelo número de nós que você quer usar no teste do NCCL:

    wget https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/refs/heads/master/gpudirect-rdma/nccl-test-a4x-max-jobset.yaml
    
    NUM_NODES=2
    sed "s|__NUM_NODES__|NUM_NODES|" nccl-test-a4x-max-jobset.yaml | kubectl apply -f -
    
  4. Verifique se todos os pods estão no estado Completed:

    kubectl get pods | grep allgather-worker
    

    O resultado será o seguinte:

    allgather-worker-0-0-g45d2   0/1     Completed   0          13m
    allgather-worker-0-1-prpvw   0/1     Completed   0          13m
    allgather-worker-0-2-qbwt5   0/1     Completed   0          13m
    
  5. Confira o resultado do teste do pod principal (nccl-test-nccl-test-0) de onde o teste é iniciado:

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

    O resultado será o seguinte:

    # nccl-tests version 2.17.6 nccl-headers=22809 nccl-library=22809
    # Collective test starting: all_gather_perf
    # nThread 1 nGpus 1 minBytes 1024 maxBytes 8589934592 step: 2(factor) warmup iters: 5 iters: 100 agg iters: 1 validation: 1 graph: 0
    #
    ...
    #                                                            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    45.49    0.02    0.02       0    45.29    0.02    0.02       0
            2048            64     float    none      -1    45.52    0.04    0.04       0    45.37    0.05    0.04       0
            4096           128     float    none      -1    46.02    0.09    0.08       0    45.83    0.09    0.08       0
            8192           256     float    none      -1    63.93    0.13    0.11       0    46.98    0.17    0.15       0
           16384           512     float    none      -1    46.51    0.35    0.31       0    47.11    0.35    0.30       0
           32768          1024     float    none      -1    66.32    0.49    0.43       0    50.73    0.65    0.57       0
           65536          2048     float    none      -1    49.89    1.31    1.15       0    50.04    1.31    1.15       0
          131072          4096     float    none      -1    54.68    2.40    2.10       0    52.38    2.50    2.19       0
          262144          8192     float    none      -1    54.66    4.80    4.20       0    54.06    4.85    4.24       0
          524288         16384     float    none      -1    66.28    7.91    6.92       0    65.75    7.97    6.98       0
         1048576         32768     float    none      -1    85.63   12.25   10.72       0    86.44   12.13   10.61       0
         2097152         65536     float    none      -1    68.33   30.69   26.86       0    72.32   29.00   25.37       0
         4194304        131072     float    none      -1    71.85   58.37   51.08       0    71.58   58.60   51.28       0
         8388608        262144     float    none      -1    83.80  100.10   87.59       0    85.73   97.85   85.62       0
        16777216        524288     float    none      -1   195.94   85.62   74.92       0   195.86   85.66   74.95       0
        33554432       1048576     float    none      -1   240.84  139.32  121.91       0   210.82  159.16  139.27       0
        67108864       2097152     float    none      -1   254.95  263.22  230.32       0   250.93  267.44  234.01       0
       134217728       4194304     float    none      -1   411.09  326.49  285.68       0   386.11  347.61  304.16       0
       268435456       8388608     float    none      -1   741.69  361.92  316.68       0   722.42  371.58  325.13       0
       536870912      16777216     float    none      -1  1358.44  395.21  345.81       0  1343.63  399.57  349.62       0
      1073741824      33554432     float    none      -1  2679.62  400.71  350.62       0  2585.68  415.26  363.36       0
      2147483648      67108864     float    none      -1  5281.54  406.60  355.78       0  5074.73  423.17  370.28       0
      4294967296     134217728     float    none      -1  10476.2  409.97  358.73       0  10027.5  428.32  374.78       0
      8589934592     268435456     float    none      -1  20853.9  411.91  360.42       0  20194.7  425.36  372.19       0
    # Out of bounds values : 0 OK
    # Avg bus bandwidth    : 126.85
    #
    # Collective test concluded: all_gather_perf