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

Linux

---

Observação: este documento descreve a configuração de um cluster que usa um tipo de máquina a3-highgpu-8g. Para implantar um cluster a3-megagpu-8g que usa GPUDirect-TCPXO, consulte Implantar um cluster A3 Mega Slurm para treinamento de ML.

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 tipo de máquina a3-highgpu-8g tem oito GPUs NVIDIA H100 conectadas, o que oferece 80 GB de memória de GPU 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, na sigla em inglês).

Ao trabalhar com VMs a3-highgpu-8g, é 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-highgpu-8g que usam o Container-Optimized OS.

Visão geral

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

1. Configure uma ou mais redes de nuvem privada virtual (VPC) e defina a configuração MTU (também conhecida como jumbo frames) como 8244.
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.