Maximiza el rendimiento de la red de GPU con GPUDirect-TCPX

Linux

---

Google Cloud diseñó la familia de máquinas optimizadas para aceleradores para ofrecer el rendimiento y la eficiencia necesarios para cargas de trabajo aceleradas por GPU, como la inteligencia artificial (IA), el aprendizaje automático (AA) y la computación de alto rendimiento (HPC).

La serie de máquinas con optimización de acelerador A3 tiene 208 CPU virtuales y hasta 1,872 GB de memoria. Cada VM A3 tiene ocho GPU NVIDIA H100 conectadas, que ofrece 80 GB de memoria GPU por GPU. Estas VMs pueden obtener hasta 1,000 Gbps de ancho de banda de red, lo que las hace ideales para modelos de lenguaje grandes basados en transformadores, bases de datos y computación de alto rendimiento (HPC).

Cuando trabajas con VMs A3, puedes usar GPUDirect-TCPX para lograr la latencia más baja posible entre aplicaciones y la red. GPUDirect-TCPX es una pila de red personalizada, remota y de acceso directo a la memoria (RDMA) que aumenta el rendimiento de la red de tus VMs A3 permitiendo que las cargas útiles de los paquetes de datos se transfieran directamente de la memoria de GPU a la interfaz de red sin pasar por la CPU ni la memoria del sistema. Las VMs de A3 pueden usar GPUDirect-TCPX combinadas con la NIC virtual de Google (gVNIC) para entregar la capacidad de procesamiento más alta entre las VMs en un clúster, en comparación con los tipos de máquina optimizados para aceleradores A2 o G2.

En este documento, se muestra cómo configurar y probar el rendimiento mejorado de la red de GPU que está disponible con GPUDirect-TCPX en las VMs A3 que usan Container-Optimized OS.

Descripción general

Para probar el rendimiento de la red con GPUDirect-TCPX, completa los pasos siguientes:

1. Configura redes MTU de marcos jumbo de nube privada virtual (VPC)
2. Crea las VMs de GPU con la imagen de Container-Optimized OS cos-105-lts.
3. En cada VM, instala los controladores de GPU.
4. En cada VM, otorga a las tarjetas de interfaz de red (NIC) acceso a la GPU.
5. Ejecuta una prueba de NCCL.

Configura redes MTU de marcos jumbo

a3-highgpu-8gLas VMs tienencinco NIC físicas. Para obtener el mejor rendimiento de estas, debes crear cinco redes de nube privada virtual y establecer la MTU en8244.

Crea una red de administración, subred y regla de firewall

Completa los siguientes pasos para configurar la red de administración:

1. Crea la red de administración mediante el comando networks create:

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

2. Crea la subred de administración mediante el 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. Crea reglas de firewall mediante el comando firewall-rules create.

   1. Crea una regla de firewall para la red de administración.

      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. Crea la regla de firewall tcp:22 para limitar las direcciones IP de origen que pueden conectarse a la VM mediante 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. Crea la regla de firewall icmp que se puede usar para verificar si hay problemas de transmisión de datos en la red.

      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
      

Reemplaza lo siguiente:

• NETWORK_NAME_PREFIX: el prefijo del nombre que se usará para las subredes y las redes de nube privada virtual.
• PROJECT_ID: El ID de tu proyecto
• REGION: la región en la que deseas crear la política.
• SSH_SOURCE_IP_RANGE: el rango de IP en formato CIDR. Esto especifica qué direcciones IP de origen pueden conectarse a tu VM mediante SSH.

Crea redes de datos, subredes y reglas de firewall

Usa el siguiente comando para crear cuatro redes de datos, cada una con subredes y reglas 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 obtener más información sobre cómo crear redes de nube privada virtual, consulta Crea y verifica una red MTU de marcos jumbo.

Crea tus VMs de GPU

Para probar el rendimiento de la red con GPUDirect-TCPX, debes crear al menos dos VMs A3.

1. Crea cada VM con la imagen de Container-Optimized OS cos-105-lts y especifica las redes MTU virtuales que se crearon en el paso anterior.

   Las VMs también deben usar la interfaz de red de NIC virtual de Google (gVNIC). Para las VMs A3, se requiere la versión 1.4.0rc3 o una posterior de gVNIC. Esta versión del controlador está disponible en Container-Optimized OS.

   La primera NIC virtual se usa como la principal para las redes y el almacenamiento generales, las otras cuatro NIC virtuales están alineadas con NUMA y dos de las ocho GPU usan el mismo interruptor de 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
   

   Reemplaza lo siguiente:

   • VM_NAME: El nombre de tu VM.
   • PROJECT_ID: El ID de tu proyecto
   • ZONE: La zona para la VM.
   • NETWORK_NAME_PREFIX: el prefijo del nombre que se usará para las subredes y las redes de nube privada virtual.

Instalar los controladores de GPU.

En cada VM A3, completa los siguientes pasos.

1. Instala los controladores de GPU de NVIDIA ejecutando el siguiente comando:

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

2. Vuelve a activar la ruta ejecutando el siguiente comando:

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

Otorga a las NIC acceso a las GPU

En cada VM de A3, otorga a las NIC acceso a las GPU mediante los siguientes pasos:

1. Configura el registro.

   • Si usas Container Registry, ejecuta el siguiente comando:

     docker-credential-gcr configure-docker
     

   • Si usas Artifact Registry, ejecuta el siguiente comando:

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

2. Configura el administrador de ruta de recepción de datos. Un servicio de administración, GPUDirect-TCPX Receive Data Path Manager, debe ejecutarse junto con las aplicaciones que usan GPUDirect-TCPX. Para iniciar el servicio en cada VM de Container-Optimized OS, ejecuta el siguiente comando:

   docker run --pull=always --rm \
     --name receive-datapath-manager \
     --detach \
     --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. Verifica que se haya iniciado el contenedor receive-datapath-manager.

   docker container logs --follow receive-datapath-manager
   

   El resultado debe parecerse al siguiente:

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

4. Para dejar de ver los registros, presiona ctrl-c.

5. Instala las reglas de las tablas de IP.

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

6. Configura la biblioteca de comunicación colectiva de NVIDIA (NCCL) y el complemento GPUDirect-TCPX.

   Se requiere una versión específica de la biblioteca NCCL y la combinación de objeto binario del complemento GPUDirect-TCPX para usar NCCL con la compatibilidad con GPUDirect-TCPX. Google Cloud proporcionó paquetes que cumplen con este requisito.

   Para instalar el paquete de Google Cloud, ejecuta el siguiente 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
   

   Si este comando se ejecuta de forma correcta, los archivos libnccl-net.so y libnccl.so se colocan en el directorio /var/lib/tcpx/lib64.

Ejecutar pruebas

En cada VM A3, ejecuta los pasos siguientes para ejecutar una prueba NCCL:

1. Inicia el contenedor.

   #!/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 \
     "$@"
   }
   

   El comando anterior completa lo siguiente:

   • Activa los dispositivos NVIDIA de /dev en el contenedor.
   • Configura el espacio de nombres de red del contenedor en el host.
   • Configura el espacio de nombres del usuario del contenedor que alojarás.
   • Agrega CAP_IPC_LOCK a las capacidades del contenedor.
   • Activa /tmp del host en /tmp del contenedor.
   • Activa la ruta de instalación del complemento NCCL y GPUDirect-TCPX NCCL en el contenedor y agrega la ruta de acceso activada a LD_LIBRARY_PATH.

2. Después de iniciar el contenedor, las aplicaciones que usan NCCL pueden ejecutarse desde el contenedor. Por ejemplo, para ejecutar la prueba run-allgather, completa los siguientes pasos:

   1. En cada VM A3, ejecuta lo siguiente:

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

   2. En una VM, ejecuta los siguientes comandos:

      1. Configura la conexión entre las VMs. Reemplaza VM-0 y VM-1 por los nombres de cada VM.

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

         Esto crea un directorio /scripts/hostfiles2 en cada VM.

      2. Ejecuta la secuencia de comandos.

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

         La secuencia de comandos run-allgather tarda unos dos minutos en ejecutarse. Al final de los registros, verás los resultados all-gather.

         Si ves la siguiente línea en tus registros NCCL, esto verifica que GPUGPU-TCPX se inicializó de forma correcta.

         NCCL INFO NET/GPUDirectTCPX ver. 3.1.1.