创建启用了 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，使它们非常适合基于 Transformer 的大型语言模型、数据库和高性能计算 (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. 设置一个或多个配置了大型 MTU 的虚拟私有云 (VPC) 网络。
  2. 使用 cos-105-lts 或更高版本的 Container-Optimized OS 映像创建 GPU 虚拟机。
  3. 在每个虚拟机上，安装 GPU 驱动程序。
  4. 在每个虚拟机上，授予网络接口卡 (NIC) 访问 GPU 的权限。
  5. 运行 NCCL 测试。

设置 VPC 网络

对于这些 VPC 网络，我们建议将最大传输单元 (MTU) 设置为更大的值。较高的 MTU 值会增加数据包大小并降低数据包标头开销，从而增加载荷数据吞吐量。

如需最大限度地提高 A3 High 机器类型的网络带宽，请将 MTU 设置为 8244 字节。 如需查看其他 GPU 机器类型的建议 MTU 设置，请参阅 GPU 机器类型的 MTU 设置。

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

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

  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=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：要用于虚拟私有云网络和子网的名称前缀。

安装 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 虚拟机上，通过完成以下步骤，让网卡能够访问 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 插件。

    使用支持 GPUDirect-TCPX 的 NCCL 需要特定的 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 脚本大约需要 2 分钟才能运行完毕。日志末尾会显示 all-gather 结果。

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

        NCCL INFO NET/GPUDirectTCPX ver. 3.1.1.