AWS Sagemaker Hyperpodにおける高性能GPUメモリ転送

現代のディープラーニングインフラストラクチャでは、多くの場合、マシン間でGPU間の膨大なデータを転送する必要があります。Perplexityでは、非連続的なGPUメモリ領域をマシン間で可能な限り高速に転送するという独自の技術的課題に直面しました。ターゲットプラットフォームであるAWS p5インスタンスは、32枚のネットワークカードを通じて3200 Gbpsという驚異的なネットワーク帯域幅を提供します。この記事では、この理論上の帯域幅の97.1%を実現するカスタム高性能ネットワーキングソリューションを構築した私たちの軌跡を共有します。

技術的課題

私たちのユースケースでは、いくつかの重要な技術的要件がありました。

• 遠隔地のGPU間における非連続的なメモリーチャンクの広帯域幅転送
• 進行中の操作を中断することなく、Kubernetesデプロイメントに対して動的にノードを追加または削除する機能
• ピアツーピア通信パターンのサポート

NVIDIAのNCCLライブラリは分散ディープラーニングの事実上の標準ですが、私たちのユースケースには理想的ではありませんでした。

• NCCLは集合通信には優れていますが、静的な「ワールド」を構築する必要があり、参加ノードを調整する際にはクラスター全体を再起動しなければなりません。
• NCCLの同期通信モデルは、私たちの非同期ワークロードに対して複雑さを増大させます。
• 最適化のために、メモリ転送パターンを直接制御したいと考えました。
• 独自のソリューションを構築することは、貴重な学習の機会となりました。

現代の高性能ネットワーク

私たちのソリューションを理解するために、まず現代の高性能ネットワークが従来のネットワークとどのように異なるのかを探ってみましょう。

日常的に使用するネットワークのほとんどはTCP/IPプロトコルに依存しており、アプリケーションはソケットを使用してオペレーティングシステムのカーネルを介しネットワークカードと通信します。しかし、高性能ネットワークはRDMA（Remote Direct Memory Access）を使用します。これは、CPUを介在させずにマシン間で直接メモリへのアクセスを可能にする、完全に異なるハードウェアおよびソフトウェアスタックです。

AWSは、Amazon独自のプロトコルであるScalable Reliable Datagram（SRD）を実装したカスタムネットワークインターフェース、Elastic Fabric Adapter（EFA）を提供しています。データがユーザー空間、カーネル空間、ネットワークバッファの間で何度もコピーされなければならない従来のTCP/IPネットワーキングとは異なり、RDMAを備えたEFAは、GPUメモリとネットワークカード間の直接的なデータ転送を可能にし、CPUを完全にバイパスします。

高性能ネットワーク設計の哲学

高性能ネットワーキングシステムを構築するには、いくつかの基本的な前提を見直す必要があります。

1. バッファの所有権：カーネルがネットワークバッファを管理し、ユーザー空間とカーネル空間の間でコピーを必要とする従来のソケットとは異なり、RDMAではアプリケーションが独自のバッファを管理する必要があります。アプリケーションがネットワーク操作を開始すると、操作が完了するまでバッファの所有権をネットワークカードに移転し、データコピーの必要性を排除します。
2. メモリ登録：アプリケーションは、オペレーティングシステムのカーネルにメモリ領域を登録する必要があります。カーネルは、CPU、GPU、およびネットワークカードがすべて同じ仮想アドレスを理解できるようにする仮想アドレスマッピングを設定します。この登録は一度限りの操作であり、その後のゼロコピーデータ転送を可能にします。
3. 制御プレーンとデータプレーン：高性能ネットワークは、操作を以下の2つのカテゴリに分類します。
   1. 制御プレーンの操作（接続設定やメモリ登録など）は、セキュリティを確保するためにカーネルを経由します。
   2. データプレーンの操作（実際のデータ転送）は、最大のパフォーマンスを得るためにカーネルをバイパスします。
4. 送信前の受信待ち受け：カーネル管理のバッファがない場合、アプリケーションは、着信データがどこに配置されるかを指定して、受信操作を事前にポストしておく必要があります。これは、アプリケーションがいつでもデータを受信できるソケットモデルとは根本的に異なります。
5. ポーリングベースの完了確認：epollのようなメカニズムを通じたカーネル通知を待つのではなく、アプリケーションがハードウェアの完了キューを直接ポーリングします。これにより、システムコールのオーバーヘッドが排除され、完了した操作に対して即座に反応することが可能になります。
6. ハードウェアトポロジへの対応：最大限のパフォーマンスを達成するには、ハードウェアトポロジを理解し、最適化することが極めて重要です。

ハードウェアトポロジの理解

AWS p5インスタンスは、高度な内部アーキテクチャを備えています。以下に示すように、各インスタンスには2つのNUMAノードを形成する2つのCPUソケットが含まれており、各NUMAノードは4つのPCIeスイッチに接続されています。

拡大

各PCIeスイッチの下には以下が存在します。

• 100 Gbps EFAネットワークカード 4枚
• NVIDIA H100 GPU 1基
• NVMe SSD 1台

TCP/IP転送とRDMA転送のデータパスは、そのアプローチにおいて根本的な違いを示しています。

拡大

TCP/IP（左側）では、データは何度もコピーされる必要があります。

1. GPUメモリからユーザー空間のアプリケーションバッファへ
2. ユーザー空間バッファからカーネルバッファへ
3. カーネルバッファからネットワークカードへ
4. 受信側でも逆のプロセスが発生します

各コピー操作はCPUサイクルとメモリ帯域幅を消費します。また、アプリケーションはネットワーク操作ごとにユーザー空間とカーネル空間の間でコンテキストスイッチを行う必要があります。

対照的に、RDMA（右側）は真のゼロコピーデータ転送を提供します。

1. ネットワークカードがGPUメモリから直接読み取る
2. データがリモートのネットワークカードへ直接送られる
3. リモートのネットワークカードが宛先のGPUメモリへ直接書き込む
4. アプリケーションは、ユーザー空間の完了キュー（CQ）を確認するだけで転送が完了したかどうかを知ることができる

RDMAと適切なハードウェアの組み合わせにより、2つのGPU間でのデータ転送には、ローカルのPCIeスイッチとネットワークを経由するだけで済みます。

拡大

対照的に、TCP/IP転送ではメインメモリを介して何度もデータをコピーする必要があり、PCIeバスの深刻な輻輳を引き起こします。

拡大

libfabricを使用した構築

私たちは、ファブリックサービス用の汎用インターフェースを提供するフレームワークであるlibfabricを使用しました。私たちの実装では、2種類のRDMA操作を使用しています。

• メモリ領域に関するメタデータ（オフセットやサイズなど）を保持する制御メッセージ用の双方向RDMA（SEND/RECV）。
• 実際のデータ転送用の片方向RDMA WRITE。各WRITE操作は1つの連続したメモリチャンクを処理します。

ピークパフォーマンスへの旅

私たちの開発はいくつかの段階を経て進行しました。

1. SEND/RECVを使用した基本的な一方向のメッセージ転送を実装
2. 双方向通信へ拡張
3. GPU-GPU間直接転送のためにGPUDirect RDMA WRITEを追加
4. 複数の同時転送を処理できるように拡張
5. 堅牢性のために操作キューイングを導入
6. 単一のネットワークカードで97.4%の帯域幅使用率を達成

32枚のネットワークカードにスケーリングする際、私たちはいくつかの重要な最適化を実装しました。

1. 操作キューイング：ネットワークカードに直接操作を送信するのではなく、アプリケーションレベルのキューを維持します。これにより、ネットワーク輻輳に対する堅牢性が提供され、プログラミングモデルが簡素化されます。
2. ネットワークのウォームアップ：事前に接続を確立することで、起動時のパフォーマンスが向上しました。
3. マルチスレッディング：各GPUのネットワーク操作専用のスレッドを用意。
4. CPUコアピンニング：スレッドを特定のCPUコアにバインドし、NUMAの影響やキャッシュミスを回避。
5. 状態のシャーディング：スレッド間の競合を削減。
6. 操作のバッチ処理：複数の操作をまとめて送信。
7. 遅延操作ポスト：操作はまずアプリケーション内でキューイングされます。完了キューをポーリングした後、保留中の操作をネットワークカードに送信するように試みることで、ネットワークリソースの効率的な使用を確実にします。
8. NUMA対応のリソース割り当て：メモリへのアクセス遅延を最小限に抑えるため、完了キューなどのlibfabricリソースを正しいNUMAノードに割り当てます。

これらの最適化を通じて、理論上の最大帯域幅の97.1%にあたる3,108 Gbpsの最終的なパフォーマンスを達成しました。

以下の動画は、私たちのコマンドラインプログラムの動作を示しています。これは非連続的なGPUメモリのチャンクをリモートノードに転送し、3108.283 Gbpsの伝送速度を達成しており、ネットワークインフラストラクチャのほぼ理論上の帯域幅使用率を実証しています。

結論

高性能ネットワーキングシステムを構築するには、ハードウェアアーキテクチャとシステム設計原則の両方を理解する必要があります。NCCLのようなライブラリは一般的なパターンに対して優れたソリューションを提供しますが、特定の要件に対してはカスタムソリューションが必要になる場合があります。

私たちの軌跡は、システムアーキテクチャ、ハードウェアトポロジ、およびさまざまな最適化技術に細心の注意を払うことで、理論上のネットワークパフォーマンスに近い性能を実現可能であることを示しています。重要なのは個々のコンポーネントを理解するだけでなく、それらがどのように相互作用して完全なシステムを形成するかを理解することです。

この旅の完全な技術的な深掘り（実装の詳細やコード例を含む）は、私たちのオープンソースリポジトリで利用可能です。

引用

@online{perplexity2025HighPerformanceGpuMemoryTransferOnAwsSagemakerHyperpod,
  author = {{Perplexity} and Lequn Chen and Phi Hoang},
  title = {{High-performance GPU memory transfer on AWS Sagemaker Hyperpod}},
  date = {2025-02-10},
  url = {https://www.perplexity.ai/hub/blog/high-performance-gpu-memory-transfer-on-aws-sagemaker-hyperpod}
}