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-0VM-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.