Criar uma VM A3 com o GPUDirect-TCPX ativado


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 tem oito GPUs NVIDIA H100 anexadas, 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 ou a3-edgegpu-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.

Este documento mostra como criar uma VM a3-highgpu-8g ou a3-edgegpu-8g que seja executada em um sistema operacional Container-Optimized OS. Ele também mostra como ativar o GPUDirect-TCPX na VM e configurar e testar o desempenho aprimorado da rede da GPU.

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 ou uma versão mais recente 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

As VMs a3-highgpu-8g e a3-edgegpu-8g têm cinco NICs físicas. Para ter o melhor desempenho das NICs físicas, crie cinco redes de nuvem privada virtual e defina a MTU como 8244.

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 do SO cos-105-lts ou mais recente do Container-Optimized OS e especifique 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=MACHINE_TYPE \
      --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: uma zona que oferece suporte ao tipo de máquina.
    • MACHINE_TYPE: o tipo de máquina da VM. Especifique a3-highgpu-8g ou a3-edgegpu-8g.
    • 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.