使用 GPUDirect-TCPX 最大限度地提高 GPU 网络性能

Linux

---

加速器优化机器系列由 Google Cloud 设计，旨在为人工智能 (AI)、机器学习 (ML) 和高性能计算 (HPC) 等 GPU 加速工作负载提供所需的性能和效率。

A3 加速器优化机器系列具有 208 个 vCPU 和高达 1872 GB 的内存。每个 A3 虚拟机都挂接了 8 个 NVIDIA H100 GPU，每个 GPU 提供 80 GB 的 GPU 内存。这些虚拟机高达 1000 Gbps 的网络带宽，因此非常适合基于转换器的大型语言模型、数据库和高性能计算 (HPC)。

使用 A3 虚拟机时，您可以使用 GPUDirect-TCPX 来实现应用和网络之间的最短延迟时间。GPUDirect-TCPX 是一种自定义远程直接内存访问 (RDMA) 网络堆栈，它允许数据包载荷从 GPU 内存直接传输到网络接口（不必通过 CPU 和系统内存），从而提升 A3 虚拟机的网络性能。与 A2 或 G2 加速器优化机器类型相比，A3 虚拟机可以将 GPUDirect-TCPX 与 Google 虚拟 NIC (gVNIC) 搭配使用，在集群中的虚拟机之间提供最高的吞吐量。

本文档介绍如何在使用和测试 Container-Optimized OS 的 A3 虚拟机上设置和测试可通过 GPU Direct-TCPX 获得的改进 GPU 网络性能。

概览

如需使用 GPUDirect-TCPX 测试网络性能，请完成以下步骤：

1. 设置 Virtual Private Cloud (VPC) MTU 巨型帧网络。
2. 使用 cos-105-lts Container-Optimized OS 映像创建 GPU 虚拟机。
3. 在每个虚拟机上，安装 GPU 驱动程序。
4. 在每个虚拟机上，向网络接口卡 (NIC) 授予 GPU 的访问权限。
5. 运行 NCCL 测试。

设置巨型帧 MTU 网络

a3-highgpu-8g 虚拟机具有五个物理 NIC，为获得最佳物理 NIC 性能，您需要创建 5 个 Virtual Private Cloud 网络并将 MTU 设置为 8244。

创建管理网络、子网和防火墙规则

完成以下步骤以设置管理网络：

1. 使用 networks create 命令创建管理网络：

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

2. 使用 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. 使用 firewall-rules create 命令创建防火墙规则。

   1. 为管理网络创建防火墙规则。

      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. 创建 tcp:22 防火墙规则以限制哪些来源 IP 地址可以使用 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. 创建 icmp 防火墙规则，用于检查网络中的数据传输问题。

      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
      

替换以下内容：

• NETWORK_NAME_PREFIX：用于 Virtual Private Cloud 网络和子网的名称前缀。
• PROJECT_ID：您的项目 ID。
• REGION：要用于创建网络的区域。
• SSH_SOURCE_IP_RANGE：采用 CIDR 格式的 IP 范围。 这用于指定哪些来源 IP 地址可以使用 SSH 连接到您的虚拟机。

创建数据网络、子网和防火墙规则

使用以下命令创建四个数据网络，每个网络都包含子网和防火墙规则。

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

如需详细了解如何创建 Virtual Private Cloud 网络，请参阅创建和验证巨型帧 MTU 网络。

创建 GPU 虚拟机

如需使用 GPUDirect-TCPX 测试网络性能，您需要创建至少两个 A3 虚拟机。

1. 通过使用 cos-105-lts Container-Optimized OS 映像并指定上一步中创建的虚拟 MTU 网络来创建每个虚拟机。

   虚拟机还必须使用 Google 虚拟 NIC (gVNIC) 网络接口。对于 A3 虚拟机，需要 gVNIC 1.4.0rc3 或更高版本。Container-Optimized OS 已提供此驱动程序版本。

   第一个虚拟 NIC 用作常规网络和存储的主要 NIC，其他四个虚拟 NIC 是与同一 PCIe 交换机上的八个 GPU 中的两个对齐的 NUMA。

   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
   

   替换以下内容：

   • VM_NAME：您的虚拟机的名称。
   • PROJECT_ID：您的项目 ID。
   • ZONE：虚拟机的可用区。
   • NETWORK_NAME_PREFIX：用于 Virtual Private Cloud 网络和子网的名称前缀。

安装 GPU 驱动程序

在每个 A3 虚拟机上，完成以下步骤。

1. 运行以下命令来安装 NVIDIA GPU 驱动程序：

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

2. 通过运行以下命令重新装载路径：

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

为 NIC 授予对 GPU 的访问权限

在每个 A3 虚拟机上，通过完成以下步骤向 NIC 授予对 GPU 的访问权限：

1. 配置注册表。

   • 如果您使用的是 Container Registry，请运行以下命令：

     docker-credential-gcr configure-docker
     

   • 如果您使用的是 Artifact Registry，请运行以下命令：

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

2. 配置接收数据路径管理器。管理服务 GPUDirect-TCPX 接收数据路径管理器需要与使用 GPUDirect-TCPX 的应用一起运行。要在每个 Container-Optimized OS 虚拟机上启动服务，请运行以下命令：

   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. 验证 receive-datapath-manager 容器是否已启动。

   docker container logs --follow receive-datapath-manager
   

   您应该会看到类似如下所示的输出：

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

4. 如需停止查看日志，请按 ctrl-c。

5. 安装 IP 表规则。

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

6. 配置 NVIDIA Collective Communications Library (NCCL) 和 GPUDirect-TCPX 插件。

   如需将 NCCL 与 GPUDirect-TCPX 支持搭配使用，您必须具有特定的 NCCL 库版本和 GPUDirect-TCPX 插件二进制文件组合。Google Cloud 提供的软件包可以满足此要求。

   如需安装 Google Cloud 软件包，请运行以下命令：

   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
   

   如果此命令成功，则 libnccl-net.so 和 libnccl.so 文件将放在 /var/lib/tcpx/lib64 目录中。

运行测试

在每个 A3 虚拟机上，完成以下步骤来运行 NCCL 测试：

1. 启动容器。

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

   上述命令会完成以下操作：

   • 将 NVIDIA 设备从 /dev 装载到容器中
   • 将容器的网络命名空间设置为主机
   • 将容器的用户命名空间设置为主机
   • 将 CAP_IPC_LOCK 添加到容器的功能
   • 将主机的 /tmp 装载到容器的 /tmp 上
   • 将 NCCL 和 GPUDirect-TCPX NCCL 插件的安装路径装载到容器中，并将装载的路径添加到 LD_LIBRARY_PATH

2. 启动容器后，使用 NCCL 的应用可以从容器内部运行。例如，如需运行 run-allgather 测试，请完成以下步骤：

   1. 在每个 A3 虚拟机上，运行以下命令：

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

   2. 在一个虚拟机上，运行以下命令：

      1. 在虚拟机之间建立连接。将 VM-0 和 VM-1 替换为每个虚拟机的名称。

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

         这将在每个虚拟机上创建一个 /scripts/hostfiles2 目录。

      2. 运行脚本。

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

         run-allgather 脚本运行大约需要两分钟。日志末尾会显示 all-gather 结果。

         如果您在 NCCL 日志中看到以下行，则说明 GPUDirect-TCPX 已成功初始化。

         NCCL INFO NET/GPUDirectTCPX ver. 3.1.1.