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

Linux

---

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

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

A3 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 VM 上で GPUDirect-TCPX で使用可能な GPU ネットワーク パフォーマンスの改善を設定してテストする方法について説明します。

概要

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

1. Virtual Private Cloud（VPC）MTU ジャンボ フレーム ネットワークを設定します。
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.