来源：examples/gcp_gpu_direct_tcpx

使用 SkyPilot 在 A3 虚拟机上启用 GCP GPUDirect-TCPX#

GPUDirect-TCPX 是一种高性能网络技术，可实现 GPU 与网络接口之间的直接通信。通过绕过 CPU 和系统内存，它显著提高了 A3 虚拟机的网络性能，尤其对于大型数据传输。

部署 a3-highgpu-8g 或 a3-edgegpu-8g 虚拟机时，将 GPUDirect-TCPX 与 Google Virtual NIC (gVNIC) 结合使用，可在应用程序和网络基础设施之间提供最佳网络性能，同时最大限度地减少延迟。

本示例演示了如何在 GCP 集群上使用 a3-highgpu-8g 虚拟机运行 NCCL 测试，并比较启用和未启用 GPUDirect-TCPX 时的性能差异。

TL;DR：使用 SkyPilot 启用 GPUDirect-TCPX#

通过在 SkyPilot YAML 中添加一个配置参数，即可在配备 a3-highgpu-8g 或 a3-edgegpu-8g 虚拟机的 GCP 集群上启用 GPUDirect-TCPX

config:
  gcp:
    enable_gpu_direct: true

设置 enable_gpu_direct: true 后，SkyPilot 会自动执行以下操作：

  1. 创建专用的网络基础设施

    • 1 个管理 VPC

    • 4 个数据 VPC

    • 每个 VPC 的相应子网

  2. 配置支持 GPUDirect-TCPX 的虚拟机

    • 使用指定的实例类型启动虚拟机

    • 使用兼容 GPUDirect-TCPX 的镜像

    • 安装必要的 GPU 驱动程序

    • 部署 GPUDirect-TCPX 接收数据路径管理器服务

    • 配置 NVIDIA 集合通信库 (NCCL) 和 GPUDirect-TCPX 插件

使用 GPUDirect-TCPX 运行 NCCL 测试#

完整的配置可在 gpu_direct_tcpx.yaml 中找到。该配置包括：

  • image_id：用于 NCCL 测试的预配置环境

  • instance_type：设置为 a3-highgpu-8g

要使用 GPUDirect-TCPX 运行 NCCL 测试

sky launch -c tcpx gpu_direct_tcpx.yaml

SkyPilot 将会：

  1. 部署一个启用 GPUDirect-TCPX 节点的 GCP 集群

  2. 执行 NCCL 性能测试

  3. 输出详细的性能指标

以下日志条目确认 GPUDirect-TCPX 已成功激活

NCCL INFO NET/GPUDirectTCPX ver. 3.1.8.
NCCL INFO NET/GPUDirectTCPX : GPUDirectTCPX enable: 1

注意： 要运行未启用 GPUDirect-TCPX 的测试，请使用

sky launch -c tcpx --env USE_GPU_DIRECT=false gpu_direct_tcpx.yaml

性能基准测试结果#

我们使用 NCCL 测试在拥有 2x a3-highgpu-8g (2xH100:8) 实例的 GCP 集群上进行了性能比较。加速比计算方式如下：

Speed-up = busbw GPUDirect-TCPX (GB/s) / busbw Non-GPUDirect-TCPX (GB/s)

消息大小

总线带宽 GPUDirect-TCPX (GB/s)

总线带宽 非 GPUDirect-TCPX (GB/s)

加速比

8 B

0

0

-

16 B

0

0

-

32 B

0

0

-

64 B

0

0

-

128 B

0

0

-

256 B

0

0

-

512 B

0

0

-

1 KB

0.01

0.01

1 倍

2 KB

0.01

0.01

1 倍

4 KB

0.01

0.02

0.5 倍

8 KB

0.02

0.04

0.5 倍

16 KB

0.04

0.09

0.4 倍

32 KB

0.09

0.12

0.7 倍

64 KB

0.11

0.17

0.6 倍

128 KB

0.19

0.15

1.2 倍

256 KB

0.35

0.23

1.5 倍

512 KB

0.65

0.47

1.4 倍

1 MB

1.33

0.95

1.4 倍

2 MB

2.43

1.87

1.3 倍

4 MB

4.8

3.64

1.3 倍

8 MB

9.21

7.1

1.3 倍

16 MB

17.16

8.83

1.9 倍

32 MB

30.08

12.07

2.5 倍

64 MB

45.31

12.48

3.6 倍

128 MB

61.58

16.27

3.8 倍

256 MB

67.82

20.93

3.2 倍

512 MB

67.09

19.93

3.3 倍

1 GB

66.2

20.09

3.3 倍

2 GB

65.72

19.39

3.4 倍

关键性能洞察#

消息大小范围

性能特征

≤ 128 KB

收益微乎其微 - GPUDirect-TCPX 对于小消息可能引入轻微开销，带宽与非 GPUDirect 模式相当或更低

256 KB – 8 MB

中等提升 - 加速比为 1.5–1.9 倍，性能交叉点在 128–256 KB

≥ 16 MB

显著优势 - 加速比为 2.5–3.8 倍，GPUDirect-TCPX 保持 65–67 GB/s，而未启用时约为 20 GB/s

GPUDirect-TCPX 的直接 GPU 到网卡通信路径消除了 CPU 和系统内存瓶颈，为大规模数据传输提供了卓越的吞吐量。这使得它对于分布式深度学习工作负载和高性能计算应用特别有效。

包含的文件#

gpu_direct_tcpx.yaml

name: nccl-gpu-direct-tcpx

resources:
  cloud: gcp
  instance_type: a3-highgpu-8g
  image_id: docker:us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx

num_nodes: 2

envs:
  USE_GPU_DIRECT: "true"

setup: |
  # Check if the /usr/local/tcpx/lib64/libnccl.so.2 is
  # present to ensure the user-data script has completed
  while [ ! -f /usr/local/tcpx/lib64/libnccl.so.2 ]; do
    echo "Waiting for user-data script to complete"
    sleep 10
  done
  # Remount the directories with exec permissions
  sudo mount -o remount,exec /usr/local/tcpx/lib64
  sudo mount -o remount,exec /usr/local/nvidia/lib64
  sudo mount -o remount,exec /usr/local/nvidia/bin


run: |
  if [ "${SKYPILOT_NODE_RANK}" == "0" ]; then
    echo "Head node"

    # Total number of processes, NP should be the total number of GPUs in the cluster
    NP=$(($SKYPILOT_NUM_GPUS_PER_NODE * $SKYPILOT_NUM_NODES))

    # Append :${SKYPILOT_NUM_GPUS_PER_NODE} to each IP as slots
    nodes=""
    for ip in $SKYPILOT_NODE_IPS; do
      nodes="${nodes}${ip}:${SKYPILOT_NUM_GPUS_PER_NODE},"
    done
    nodes=${nodes::-1}
    echo "All nodes: ${nodes}"

    # Set environment variables
    export PATH=/usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
    export NCCL_SOCKET_IFNAME=eth0
    export NCCL_CROSS_NIC=0
    export NCCL_ALGO=Ring
    export NCCL_PROTO=Simple
    export NCCL_NSOCKS_PERTHREAD=4
    export NCCL_SOCKET_NTHREADS=1
    export NCCL_NET_GDR_LEVEL=PIX
    export NCCL_DYNAMIC_CHUNK_SIZE=524288
    export NCCL_P2P_PXN_LEVEL=0    
    export NCCL_P2P_NET_CHUNKSIZE=524288
    export NCCL_P2P_PCI_CHUNKSIZE=524288
    export NCCL_P2P_NVL_CHUNKSIZE=1048576
    export NCCL_BUFFSIZE=8388608
    export NCCL_MAX_NCHANNELS=8
    export NCCL_MIN_NCHANNELS=8
    export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
    export NCCL_GPUDIRECTTCPX_SOCKET_IFNAME=eth1,eth2,eth3,eth4
    export NCCL_GPUDIRECTTCPX_CTRL_DEV=eth0
    export NCCL_GPUDIRECTTCPX_TX_BINDINGS="eth1:8-21,112-125;eth2:8-21,112-125;eth3:60-73,164-177;eth4:60-73,164-177"
    export NCCL_GPUDIRECTTCPX_RX_BINDINGS="eth1:22-35,126-139;eth2:22-35,126-139;eth3:74-87,178-191;eth4:74-87,178-191"
    export NCCL_GPUDIRECTTCPX_PROGRAM_FLOW_STEERING_WAIT_MICROS=50000
    export NCCL_GPUDIRECTTCPX_UNIX_CLIENT_PREFIX="/run/tcpx"
    export NCCL_GPUDIRECTTCPX_FORCE_ACK=0

    if [ "${USE_GPU_DIRECT}" == "true" ]; then
      export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/tcpx/lib64
    else
      # Use the default NCCL library
      export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64
    fi
    
    mpirun \
      --allow-run-as-root \
      --tag-output \
      -H $nodes \
      -np $NP \
      -N $SKYPILOT_NUM_GPUS_PER_NODE \
      -x PATH \
      -x LD_LIBRARY_PATH \
      -x NCCL_SOCKET_IFNAME \
      -x NCCL_CROSS_NIC \
      -x NCCL_ALGO \
      -x NCCL_PROTO \
      -x NCCL_NSOCKS_PERTHREAD \
      -x NCCL_SOCKET_NTHREADS \
      -x NCCL_MAX_NCHANNELS \
      -x NCCL_MIN_NCHANNELS \
      -x NCCL_DYNAMIC_CHUNK_SIZE \
      -x NCCL_P2P_NET_CHUNKSIZE \
      -x NCCL_P2P_PCI_CHUNKSIZE \
      -x NCCL_P2P_NVL_CHUNKSIZE \
      -x NCCL_BUFFSIZE \
      -x CUDA_VISIBLE_DEVICES \
      -x NCCL_GPUDIRECTTCPX_SOCKET_IFNAME \
      -x NCCL_GPUDIRECTTCPX_CTRL_DEV \
      -x NCCL_GPUDIRECTTCPX_TX_BINDINGS \
      -x NCCL_GPUDIRECTTCPX_RX_BINDINGS \
      -x NCCL_GPUDIRECTTCPX_PROGRAM_FLOW_STEERING_WAIT_MICROS \
      -x NCCL_GPUDIRECTTCPX_UNIX_CLIENT_PREFIX \
      -x NCCL_GPUDIRECTTCPX_FORCE_ACK \
      -x NCCL_NET_GDR_LEVEL \
      -x NCCL_P2P_PXN_LEVEL \
      -x NCCL_DEBUG=INFO -x NCCL_DEBUG_SUBSYS=ENV \
      --mca btl tcp,self \
      --mca btl_tcp_if_include eth0 \
      --mca plm_rsh_args "-p 10022" \
      /third_party/nccl-tests-mpi/build/all_reduce_perf \
      -b 8 \
      -e 2G \
      -f 2 \
      -g 1 \
      -c 1 \
      -w 5 \
      -n 20
  else
    echo "Worker nodes"
  fi

config:
  gcp:
    enable_gpu_direct: true
    managed_instance_group:
      run_duration: 36000
      provision_timeout: 900