Maximizar o desempenho da rede da GPU com o GPUDirect-TCPX

A família de máquinas com otimização para aceleradores foi projetada pelo Google Cloud para oferecer o desempenho e a eficiência necessários às cargas de trabalho aceleradas por GPU, como inteligência artificial (IA), machine learning (ML) e computação de alto desempenho (HPC).

A série de máquinas com otimização para aceleradores A3 tem 208 vCPUs e até 1872 GB de memória. Cada VM A3 tem oito GPUs NVIDIA H100 anexadas, que oferecem 80 GB de memória por GPU. Essas VMs podem ter até 1.000 Gbps de largura de banda de rede, o que as torna ideais para grandes modelos de linguagem baseados em transformadores, bancos de dados e computação de alto desempenho (HPC).

Ao trabalhar com VMs do A3, é possível usar o GPUDirect-TCPX para alcançar a menor latência possível entre aplicativos e a rede. GPUDirect-TCPX é uma pilha de rede personalizada e remota de acesso direto à memória (RDMA, na sigla em inglês) que aumenta o desempenho da rede das VMs A3, permitindo que os payloads de pacotes de dados sejam transferidos diretamente da memória GPU para a interface de rede sem precisar passar pela CPU e pela memória do sistema. As VMs A3 podem usar o GPUDirect-TCPX combinado com a placa de rede virtual do Google (gVNIC) para oferecer a maior capacidade entre as VMs em um cluster, em comparação com os tipos de máquina otimizados por acelerador A2 ou G2.

Neste documento, mostramos como configurar e testar o desempenho aprimorado da rede GPU que está disponível com o GPUDirect-TCPX em VMs A3 que usam o Container-Optimized OS.

Informações gerais

Para testar o desempenho da rede com o GPUDirect-TCPX, siga estas etapas:

  1. Configurar redes de frames jumbo de MTU de nuvem privada virtual (VPC).
  2. Crie suas VMs de GPU usando a imagem cos-105-lts do Container-Optimized OS.
  3. Em cada VM, instale os drivers da GPU.
  4. Em cada VM, conceda acesso à GPU às placas de interface de rede (NICs).
  5. Execute um teste da NCCL.

Configurar redes MTU jumbo frame

a3-highgpu-8gAs VMs têmcinco NICs físicas, para ter o melhor desempenho das NICs físicas, crie cinco redes de nuvem privada virtual e defina a MTU como8244 de dois minutos.

Criar rede de gerenciamento, sub-rede e regra de firewall

Siga estas etapas para configurar a rede de gerenciamento:

  1. Crie a rede de gerenciamento usando o comando networks create:

    gcloud compute networks create NETWORK_NAME_PREFIX-mgmt-net \
  --project=PROJECT_ID \
  --subnet-mode=custom \
  --mtu=8244

  2. Crie a sub-rede de gerenciamento usando o comando networks subnets create:

    gcloud compute networks subnets create NETWORK_NAME_PREFIX-mgmt-sub \
  --project=PROJECT_ID \
  --network=NETWORK_NAME_PREFIX-mgmt-net \
  --region=REGION \
  --range=192.168.0.0/24

  3. Crie regras de firewall usando o comando firewall-rules create.

    1. Crie uma regra de firewall para a rede de gerenciamento.

      gcloud compute firewall-rules create NETWORK_NAME_PREFIX-mgmt-internal \
 --project=PROJECT_ID \
 --network=NETWORK_NAME_PREFIX-mgmt-net \
 --action=ALLOW \
 --rules=tcp:0-65535,udp:0-65535,icmp \
 --source-ranges=192.168.0.0/16

    2. Crie a regra de firewall tcp:22 para limitar quais endereços IP de origem podem se conectar à sua VM usando SSH.

      gcloud compute firewall-rules create NETWORK_NAME_PREFIX-mgmt-external-ssh \
 --project=PROJECT_ID \
 --network=NETWORK_NAME_PREFIX-mgmt-net \
 --action=ALLOW \
 --rules=tcp:22 \
 --source-ranges=SSH_SOURCE_IP_RANGE

    3. Crie a regra de firewall icmp que pode ser usada para verificar se há problemas de transmissão de dados na rede.

      gcloud compute firewall-rules create NETWORK_NAME_PREFIX-mgmt-external-ping \
 --project=PROJECT_ID \
 --network=NETWORK_NAME_PREFIX-mgmt-net \
 --action=ALLOW \
 --rules=icmp \
 --source-ranges=0.0.0.0/0

Substitua:

  • NETWORK_NAME_PREFIX: o prefixo de nome a ser usado para as redes e sub-redes de nuvem privada virtual.
  • PROJECT_ID: o ID do projeto.
  • REGION: a região em que você quer criar a política.
  • SSH_SOURCE_IP_RANGE: intervalo de IP no formato CIDR. Isso especifica quais endereços IP de origem podem se conectar à VM usando SSH.

Criar redes de dados, sub-redes e regra de firewall

Use o comando a seguir para criar quatro redes de dados, cada uma com sub-redes e regras de firewall.

for N in $(seq 1 4); do
  gcloud compute networks create NETWORK_NAME_PREFIX-data-net-$N \
      --project=PROJECT_ID \
      --subnet-mode=custom \
      --mtu=8244

  gcloud compute networks subnets create NETWORK_NAME_PREFIX-data-sub-$N \
      --project=PROJECT_ID \
      --network=NETWORK_NAME_PREFIX-data-net-$N \
      --region=REGION \
      --range=192.168.$N.0/24

  gcloud compute firewall-rules create NETWORK_NAME_PREFIX-data-internal-$N \
      --project=PROJECT_ID \
      --network=NETWORK_NAME_PREFIX-data-net-$N \
      --action=ALLOW \
      --rules=tcp:0-65535,udp:0-65535,icmp \
      --source-ranges=192.168.0.0/16
done

Para mais informações sobre como criar redes de nuvem privada virtual, consulte Criar e verificar uma rede MTU de frame jumbo.

Criar VMs de GPU

Para testar o desempenho da rede com o GPUDirect-TCPX, é necessário criar pelo menos duas VMs A3.

  1. Crie cada VM usando a imagem cos-105-lts do Container-Optimized OS e especificando as redes MTU virtuais que foram criadas na etapa anterior.

    As VMs também precisam usar a interface de rede NIC virtual do Google (gVNIC). Para VMs do A3, é necessária a versão 1.4.0rc3 ou posterior da gVNIC. Essa versão do driver está disponível no Container-Optimized OS.

    A primeira NIC virtual é usada como principal para rede e armazenamento gerais. As outras quatro NICs virtuais são alinhadas a duas das oito GPUs no mesmo switch PCIe.

    gcloud compute instances create VM_NAME \
  --project=PROJECT_ID \
  --zone=ZONE \
  --machine-type=a3-highgpu-8g \
  --maintenance-policy=TERMINATE --restart-on-failure \
  --image-family=cos-105-lts \
  --image-project=cos-cloud \
  --boot-disk-size=${BOOT_DISK_SZ:-50} \
  --metadata=cos-update-strategy=update_disabled \
  --scopes=https://www.googleapis.com/auth/cloud-platform \
  --network-interface=nic-type=GVNIC,network=NETWORK_NAME_PREFIX-mgmt-net,subnet=NETWORK_NAME_PREFIX-mgmt-sub \
  --network-interface=nic-type=GVNIC,network=NETWORK_NAME_PREFIX-data-net-1,subnet=NETWORK_NAME_PREFIX-data-sub-1,no-address \
  --network-interface=nic-type=GVNIC,network=NETWORK_NAME_PREFIX-data-net-2,subnet=NETWORK_NAME_PREFIX-data-sub-2,no-address \
  --network-interface=nic-type=GVNIC,network=NETWORK_NAME_PREFIX-data-net-3,subnet=NETWORK_NAME_PREFIX-data-sub-3,no-address \
  --network-interface=nic-type=GVNIC,network=NETWORK_NAME_PREFIX-data-net-4,subnet=NETWORK_NAME_PREFIX-data-sub-4,no-address

    Substitua:

    • VM_NAME: o nome da VM.
    • PROJECT_ID: o ID do projeto.
    • ZONE: a zona para a VM.
    • NETWORK_NAME_PREFIX: o prefixo de nome a ser usado para as redes e sub-redes da nuvem privada virtual.

Instalar os drivers da GPU

Em cada VM A3, conclua as etapas a seguir.

  1. Instale os drivers de GPU NVIDIA executando o seguinte comando:

    sudo cos-extensions install gpu -- --version=latest

  2. Remonte o caminho executando o seguinte comando:

    sudo mount --bind /var/lib/nvidia /var/lib/nvidia
sudo mount -o remount,exec /var/lib/nvidia

Conceder às placas de rede (NICs) acesso às GPUs

Em cada VM A3, conceda às placas de rede acesso às GPUs, concluindo as seguintes etapas:

  1. Configure o registro.

    • Se você estiver usando o Container Registry, execute o seguinte comando:

      docker-credential-gcr configure-docker

    • Se você estiver usando o Artifact Registry, execute o seguinte comando:

      docker-credential-gcr configure-docker --registries us-docker.pkg.dev

  2. Configurar o gerenciador de caminho dos dados de recebimento. Um serviço de gerenciamento, o gerenciador de caminho de dados de recebimento GPUDirect-TCPX, precisa ser executado junto com os aplicativos que usam GPUDirect-TCPX. Para iniciar o serviço em cada VM do Container-Optimized OS, execute o seguinte comando:

    docker run --pull=always --rm \
  --name receive-datapath-manager \
  --detach \
  --privileged \
  --cap-add=NET_ADMIN --network=host \
  --volume /var/lib/nvidia/lib64:/usr/local/nvidia/lib64 \
  --device /dev/nvidia0:/dev/nvidia0 \
  --device /dev/nvidia1:/dev/nvidia1 \
  --device /dev/nvidia2:/dev/nvidia2 \
  --device /dev/nvidia3:/dev/nvidia3 \
  --device /dev/nvidia4:/dev/nvidia4 \
  --device /dev/nvidia5:/dev/nvidia5 \
  --device /dev/nvidia6:/dev/nvidia6 \
  --device /dev/nvidia7:/dev/nvidia7 \
  --device /dev/nvidia-uvm:/dev/nvidia-uvm \
  --device /dev/nvidiactl:/dev/nvidiactl \
  --env LD_LIBRARY_PATH=/usr/local/nvidia/lib64 \
  --volume /run/tcpx:/run/tcpx \
  --entrypoint /tcpgpudmarxd/build/app/tcpgpudmarxd \
us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/tcpgpudmarxd \
  --gpu_nic_preset a3vm --gpu_shmem_type fd --uds_path "/run/tcpx" --setup_param "--verbose 128 2 0"

  3. Verifique se o contêiner receive-datapath-manager foi iniciado.

    docker container logs --follow receive-datapath-manager

    A saída será semelhante a esta:

    I0000 00:00:1687813309.406064       1 rx_rule_manager.cc:174] Rx Rule Manager server(s) started...

  4. Para interromper a visualização dos registros, pressione ctrl-c.

  5. Instalar as regras da tabela de IP.

    sudo iptables -I INPUT -p tcp -m tcp -j ACCEPT

  6. Configure a NVIDIA Collective Communications Library (NCCL) e o plug-in GPUDirect-TCPX.

    Uma versão específica da biblioteca NCCL e uma combinação binária do plug-in GPUDirect-TCPX são necessárias para usar o NCCL com suporte a GPUDirect-TCPX. O Google Cloud oferece pacotes que atendem a esse requisito.

    Para instalar o pacote do Google Cloud, execute o seguinte comando:

    docker run --rm -v /var/lib:/var/lib us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx install --install-nccl
sudo mount --bind /var/lib/tcpx /var/lib/tcpx
sudo mount -o remount,exec /var/lib/tcpx

    Se esse comando for bem-sucedido, os arquivos libnccl-net.so e libnccl.so serão colocados no diretório /var/lib/tcpx/lib64.

Executar testes

Em cada VM A3, execute um teste de NCCL seguindo estas etapas:

  1. Inicie o contêiner.

    #!/bin/bash

function run_tcpx_container() {
docker run \
  -u 0 --network=host \
  --cap-add=IPC_LOCK \
  --userns=host \
  --volume /run/tcpx:/tmp \
  --volume /var/lib/nvidia/lib64:/usr/local/nvidia/lib64 \
  --volume /var/lib/tcpx/lib64:/usr/local/tcpx/lib64 \
  --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 \
  --device /dev/nvidia0:/dev/nvidia0 \
  --device /dev/nvidia1:/dev/nvidia1 \
  --device /dev/nvidia2:/dev/nvidia2 \
  --device /dev/nvidia3:/dev/nvidia3 \
  --device /dev/nvidia4:/dev/nvidia4 \
  --device /dev/nvidia5:/dev/nvidia5 \
  --device /dev/nvidia6:/dev/nvidia6 \
  --device /dev/nvidia7:/dev/nvidia7 \
  --device /dev/nvidia-uvm:/dev/nvidia-uvm \
  --device /dev/nvidiactl:/dev/nvidiactl \
  --env LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/tcpx/lib64 \
  "$@"
}

    O comando anterior completa o seguinte:

    • Monta dispositivos NVIDIA do /dev no contêiner
    • Define o namespace da rede do contêiner para o host
    • Define o namespace do usuário do contêiner para host
    • Adiciona CAP_IPC_LOCK aos recursos do contêiner
    • Monta o /tmp do host no /tmp do contêiner
    • Monta o caminho de instalação do plug-in NCCL e GPUDirect-TCPX NCCL no contêiner e adiciona o caminho montado a LD_LIBRARY_PATH

  2. Depois de iniciar o contêiner, os aplicativos que usam NCCL podem ser executados dentro do contêiner. Por exemplo, para executar o teste run-allgather, conclua as seguintes etapas:

    1. Em cada VM A3, execute o seguinte:

      $ run_tcpx_container -it --rm us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx shell

    2. Em uma VM, execute os seguintes comandos:

      1. Configure a conexão entre as VMs. Substitua VM-0 e VM-1 pelos nomes de cada VM.

        /scripts/init_ssh.sh VM-0 VM-1
pushd /scripts && /scripts/gen_hostfiles.sh VM-0 VM-1; popd

        Isso cria um diretório /scripts/hostfiles2 em cada VM.

      2. Execute o script.

        /scripts/run-allgather.sh 8 eth1,eth2,eth3,eth4 1M 512M 2

        O script run-allgather leva cerca de dois minutos para ser executado. No final dos registros, você verá os resultados all-gather.

        Se a linha a seguir for exibida nos registros da NCCL, isso vai verificar se o GPUDirect-TCPX foi inicializado com sucesso.

        NCCL INFO NET/GPUDirectTCPX ver. 3.1.1.