GPUDirect-TCPX を使用して GPU ネットワークのパフォーマンスを最大化する

アクセラレータ最適化マシン ファミリーは、Google Cloud が設計し、AI、ML、ハイ パフォーマンス コンピューティング（HPC）などの GPU 高速化ワークロードに必要なパフォーマンスと効率性を実現しています。

A3 アクセラレータ最適化マシンシリーズは、208 個の vCPU と最大 1,872 GB のメモリを備えています。各 a3-highgpu-8g マシンタイプには 8 個の NVIDIA H100 GPU がアタッチされており、GPU あたり 80 GB の GPU メモリを提供します。これらの VM は最大 1,000 Gbps のネットワーク帯域幅を利用できるため、Transformer ベースの大規模言語モデル、データベース、ハイ パフォーマンス コンピューティング（HPC）に最適です。

a3-highgpu-8g VM を使用する場合は、GPUDirect-TCPX を使用することで、アプリケーションとネットワーク間のレイテンシを可能な限り低減できます。GPUDirect-TCPX は、カスタムのリモート ダイレクト メモリ アクセス（RDMA）ネットワーク スタックです。これは、CPU とシステムメモリを経由せずに、データパケットのペイロードを GPU メモリからネットワーク インターフェースに直接転送できるため、A3 VM のネットワーク パフォーマンスが向上します。A3 VM は、GPUDirect-TCPX と Google Virtual NIC（gVNIC）の組み合わせを使用して、A2 または G2 アクセラレータ最適化マシンタイプと比較して、クラスタ内の VM 間で最高のスループットを提供できます。

このドキュメントでは、Container-Optimized OS を使用する a3-highgpu-8g VM で GPUDirect-TCPX を使用して利用可能な GPU ネットワーク パフォーマンスの改善を設定してテストする方法について説明します。

概要

GPUDirect-TCPX を使用してネットワーク パフォーマンスをテストするには、次のように操作します。

  1. 1 つ以上の Virtual Private Cloud（VPC）ネットワークを設定し、MTU 設定（ジャンボ フレーム）を 8244 に設定します。
  2. cos-105-lts Container-Optimized OS イメージを使用して GPU VM を作成します。
  3. VM ごとに GPU ドライバをインストールします。
  4. 各 VM で、ネットワーク インターフェース カード（NIC）に GPU へのアクセス権を付与します。
  5. NCCL テストを実行します。

ジャンボ フレーム MTU ネットワークを設定する

a3-highgpu-8g VM には 5 つの物理 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 ファイアウォール ルールを作成して、SSH を使用して VM に接続できる送信元 IP アドレスを制限します。

      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 範囲。SSH を使用して VM に接続できる送信元 IP アドレスを指定します。

データ ネットワーク、サブネット、ファイアウォール ルールを作成する

次のコマンドを使用して、それぞれサブネットとファイアウォール ルールを持つ 4 つのデータ ネットワークを作成します。

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 VM を作成する

GPUDirect-TCPX でネットワーク パフォーマンスをテストするには、少なくとも 2 つの A3 VM を作成する必要があります。

  1. cos-105-lts Container Optimized OS イメージを使用して各 VM を作成し、前の手順で作成した仮想 MTU ネットワークを指定します。

    VM は Google Virtual NIC（gVNIC）ネットワーク インターフェースも使用する必要があります。A3 VM の場合、gVNIC バージョン 1.4.0rc3 以降が必要です。このドライバ バージョンは Container-Optimized OS で利用できます。

    最初の仮想 NIC は一般的なネットワーキングとストレージのプライマリ NIC として使用されます。他の 4 つの仮想 NIC は、同じ PCIe スイッチ上の 8 つの GPU のうち 2 つと 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: VM の名前。
    • PROJECT_ID: プロジェクト ID。
    • ZONE: VM のゾーン。
    • NETWORK_NAME_PREFIX: Virtual Private Cloud ネットワークとサブネットに使用する名前の接頭辞。

GPU ドライバをインストールする

各 A3 VM で次の手順で操作します。

  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 VM で、次の手順で 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 VM でサービスを開始するには、次のコマンドを実行します。

    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 VM で、次の手順で 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 VM で次のコマンドを実行します。

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

    2. 1 つの VM で次のコマンドを実行します。

      1. VM 間の接続を設定します。VM-0 と VM-1 は、各 VM の名前に置き換えます。

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

        これにより、各 VM に /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.