创建启用了 GPUDirect-TCPX 的 A3 虚拟机

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

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

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

本文档介绍了如何创建在 Container-Optimized OS 操作系统上运行的 a3-highgpu-8g 或 a3-edgegpu-8g 虚拟机。该文档还介绍了如何在虚拟机上启用 GPUDirect-TCPX，以及如何设置和测试经过改进的 GPU 网络性能。

概览

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

  1. 设置一个或多个虚拟私有云 (VPC) 网络，并将 MTU 设置（也称为巨帧）设为 8244。
  2. 使用 cos-105-lts 或更高版本的 Container-Optimized OS 映像创建 GPU 虚拟机。
  3. 在每个虚拟机上，安装 GPU 驱动程序。
  4. 在每个虚拟机上，向网络接口卡 (NIC) 授予对 GPU 的访问权限。
  5. 运行 NCCL 测试。

设置巨型帧 MTU 网络

a3-highgpu-8g 和 a3-edgegpu-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：要为虚拟私有云网络和子网使用的名称前缀。
  • PROJECT_ID：您的项目 ID。
  • REGION：要用于创建网络的区域。
  • SSH_SOURCE_IP_RANGE：IP 地址范围（采用 CIDR 格式）。这会指定哪些来源 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

如需详细了解如何创建虚拟私有云网络，请参阅创建和验证巨型帧 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=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

    替换以下内容：

    • VM_NAME：您的虚拟机的名称。
    • PROJECT_ID：您的项目 ID。
    • ZONE：支持您的机器类型的可用区。
    • MACHINE_TYPE：虚拟机的机器类型。指定 a3-highgpu-8g 或 a3-edgegpu-8g。
    • 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 Receive Data Path Manager 需要与使用 GPUDirect-TCPX 的应用一起运行。如需在每个 Container-Optimized OS 虚拟机上启动该服务，请运行以下命令：

    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. 验证 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 插件二进制文件组合，才能使用支持 GPUDirect-TCPX 的 NCCL。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.