著者：アイ・バオレ

導入

最近、Alibaba Cloud Machine Learning Platform (PAI) チームと DAMO Academy の GraphScope チームが共同で、PyTorch 向けの GPU アクセラレーション分散 GNN フレームワークであるGraphLearn-for-PyTorch (GLT) をリリースしました。GLT は、 GPU の強力な並列コンピューティング機能を活用してグラフのサンプリングを高速化し、UVA を利用して頂点とエッジの特徴の変換とコピーを削減します。大規模グラフの場合、GLT はプロデューサー/コンシューマーアーキテクチャを採用し、非同期の同時分散サンプリングと特徴検索、およびホットスポットキャッシュを通じて、複数の GPU またはマシンにわたる効率的な分散トレーニングをサポートします。インターフェースの面では、GLT は PyTorch スタイルを維持し、PyG と互換性があるため、わずかなコード変更のみで単一マシンの PyG トレーニングプログラムを高速化したり、単一マシンの PyG モデルを分散トレーニング用に変換したりできます。さらに、GLT は多様なニーズに対応するための柔軟な分散トレーニング展開を提供します。

オープンソースアドレス: https://github.com/alibaba/graphlearn-for-pytorch

ドキュメントアドレス: https://graphlearn-torch.readthedocs.io/en/latest/index.html

背景紹介

グラフニューラルネットワーク（Graph Neural Networks）は、グラフデータ上の表現を学習する手法として、様々なグラフ関連分野に広く応用されており、電子商取引のレコメンデーション、セキュリティリスク管理、生体分子などの分野で実用化が実現されています。グラフニューラルネットワークは、その独自のデータ処理とニューラルネットワーク計算ロジックのため、トレーニングをサポートする専用の学習フレームワークを必要とします。PAIチームは以前、大規模な産業グレードの分散グラフ学習フレームワークGraphLearn（ https://github.com/alibaba/graph-learn ）をオープンソース化しました。 GraphLearnは主にTensorFlow 1.xシリーズを使用し、psアーキテクチャによる非同期トレーニングモードを採用することで、数十億ノード、数百億エッジを持つ大規模な異種グラフの分散トレーニングをサポートします。Alibaba内外の様々なビジネスシナリオに適用されています。PyTorchの台頭により、より柔軟でPythonライクなインターフェースとデバッグの容易さから、アルゴリズム開発者の間でPyTorchが好まれる選択肢となっています。DGLやPyGといったPyTorchベースのオープンソースGNNフレームワークは、主に単一マシンベースであり、大規模グラフの分散トレーニングをサポートできません。

さらに、GPU並列コンピューティングの利点により、GPUを用いたグラフニューラルネットワークの学習は、CPU学習と比較して数倍のパフォーマンス向上を実現します。しかし、一般的なグラフニューラルネットワークフレームワークは、グラフのトポロジデータと特徴量をメモリに保存し、CPUでサンプリングと特徴量の検索を行い、その後、ニューラルネットワークの学習のためにデータをGPUにコピーします。このプロセスにおいて、グラフのサンプリングと特徴量の検索部分が学習全体のボトルネックになりやすいです。以下では、大規模グラフにおける典型的な学習フローを例に、学習プロセスのパフォーマンスボトルネックを分析し、説明します。

典型的なGNNトレーニングプロセス[1]には以下が含まれます。

1. サブグラフ トポロジ サンプリングでは、マルチホップ ネイバーをサンプリングし、サブグラフを形成します。

2. サブグラフ内のノードまたはエッジの特性を照会します。

3. サブグラフの形式をニューラル ネットワークのトレーニングに必要な形式に変換し、GPU メモリにコピーします。

4. 個別の特徴に対する埋め込み検索を実行するなど、元の特徴を処理します。

5. 近隣集約;

6. ノードの更新。

ステップ3と4はオプションです。一般的なGNNモデルはニューラルネットワークパラメータが比較的小さいため、計算コストが低く抑えられます。ボトルネックとなるのは通常、ステップ1～4、つまり通信、データのコピー、フォーマット変換などのI/O操作です。そのため、GPUをトレーニングに使用してもGPUリソ​​ースの利用率が低くなり、全体的なスループットとスケーラビリティの向上が困難になります。

要約すると、GPU などのハードウェア リソースを最大限に活用し、グラフのデータ分散特性に基づいて最適化し、さまざまなアルゴリズム モデルと並列戦略を組み合わせて通信とデータ変換時間を短縮し、全体的なスループットを向上できる、PyTorch に基づく効率的な分散 GNN トレーニング フレームワークが必要です。

キーデザイン

設計初期段階では、Quiver[2]チームと協力し、GPUサンプリングの実現可能性に関する予備調査を実施しました。その結果、GPUサンプリングはCPUと比較して桁違いのパフォーマンス向上をもたらすことが分かりました。さらに、CPUメモリとGPUメモリ間での特徴量コピーのオーバーヘッドを削減するため、特徴量をGPUメモリに直接保存することも可能です。しかし、大規模データの場合、メモリにすべての特徴量データを保存することはできません。そこで、特徴量の保存にはUVA技術を採用しました。UVA技術も、直接メモリ保存と比較して桁違いのパフォーマンス向上を実現します。業界におけるグラフの規模は、単一マシンの限界を容易に超える可能性があるため、分散トレーニングフレームワークをさらに設計しました。具体的には、プロデューサー・コンシューマーパラダイムを用いて分散グラフサンプリング、特徴量検索、モデルトレーニングの関係を処理し、グラフサンプリング、特徴量検索、トレーニングを分離し、マルチプロセス並列処理とコルーチン非同期並行処理を用いてサンプリングと特徴量検索を高速化し、プリフェッチとホットスポットキャッシュを用いてトレーニング側の待機時間をさらに短縮し、エンドツーエンドのスループットを向上させました。ユーザーの移行コストと使いやすさを考慮し、インターフェースに関してはPyGとの互換性を維持しました。PyGのトレーニングを高速化したり、GLTの分散トレーニングに移行したりするには、PyGコードにわずかな変更を加えるだけで済みます。以下に、いくつかの重要な設計ポイントを詳しく説明します。

GPUサンプリング

GLTは、グラフトポロジをCSR形式でGPUメモリまたはピンメモリに保存し、CUDAサンプリング演算子を実装してGPU並列サンプリングを実現します。CSRストレージを使用することで、各ノードの近傍ノードを容易に取得し、各ノードを個別にサンプリングできるため、GPUマルチスレッドによる並列サンプリングが容易になります。非置換ランダムサンプリングには、リザーバアルゴリズムを使用しました。バッチサイズが大きい場合、GPUサンプリングはCPUサンプリングと比較して桁違いの高速化を実現します。

ユニファイドテンソル

CPU メモリと GPU メモリ間のコピーのオーバーヘッドをなくすには、特徴量を GPU メモリに保存するのが最も簡単な方法です。しかし、単一の GPU のメモリは限られているため、特徴量データが大きい場合、すべての特徴量を GPU メモリに保存することは困難です。そのため、GLT では、べき乗分布やサンプリング アクセス特性 (次数の高いノードはより頻繁にアクセスされるなど) などのグラフ特性を利用して、一部のホット特徴量を GPU メモリに保存し、他の特徴量を RAM に保存します。GPU が RAM 内の特徴量にアクセスできるようにするために、UVA (Unified Universal Access) が使用されます。GLT は UnifiedTensor を使用して CUDA Tensor と CPU Tensor を統一的に管理し、簡潔で効率的なデータ アクセスを提供します。さらに、GPU がピアツーピア アクセス (NVLink を使用) を直接実行できる場合、GPU メモリも統一的に管理できるため、GPU メモリ内の特徴量の保存容量を拡張できます。 GLTはUnifiedTensorを用いて、これらの異なるハードウェアデバイス上のストレージを統一的に管理し、CUDAテンソルへの直接アクセス、NVLinkを介した他のGPU上のCUDAテンソルへのアクセス、そしてUVAを介したCPUテンソルへのゼロコピーアクセスを可能にします。上位層の要素ルックアップインターフェースは通常のテンソルと同様であり、下位層は対応するデバイス上でメモリアクセス操作を自動的に実行します。

特徴

機能は UnifiedTensor で構成され、ハードウェア トポロジを認識します。具体的には、まず、ユーザーが指定した CPU/GPU メモリ サイズに応じて、GPU (ホット) 部分と CPU (コールド) 部分に分割されます。次に、GPU 部分では、各カード間および NVLink 接続された各 GPU グループ間のレプリケーションなど、ユーザーが指定したレプリケーション戦略に従ってレプリケーションが実行されます。カード間レプリケーションと比較して、GPU グループ レプリケーションでは、GPU グループ内の GPU がピアツーピアで相互にアクセスできるため、より多くのホット データを GPU 上に配置できます。実装では、GLT は DeviceGroup を抽象化して、カード間レプリケーションとグループ間レプリケーションを統一的に表現します。DeviceGroup は、NVLink 接続された GPU のグループを表します。NVLink のないカードが 8 枚あると仮定すると、8 つの DeviceGroup が存在します。GPU 0 ～ 3 が NVLink を介してペアで接続され、GPU 4 ～ 7 が NVLink を介してペアで接続されている場合、GPU 0 ～ 3 が 1 つの DeviceGroup を形成し、GPU 4 ～ 7 が別の DeviceGroup を形成します。実際のテストでは、UnifiedTensor 機能を使用した場合のパフォーマンスは、CPU Tensor ルックアップ (GPU へのコピーを含む) よりも桁違いに高速であり、GPU ストレージの割合を制御することで、速度とメモリ使用量のバランスを柔軟に実現できます。

分散設計

GLT分散GNNトレーニングは、主に分散サンプリング、特徴検索、モデルトレーニングの3つの部分に分かれています。1回のサンプリング操作の結果は通常小さく（最大数十MB）、特徴検索の結果は大きくなります（数百MB）。特徴検索の結果は、トレーニング中のニューラルネットワーク計算に使用されます。特徴検索では、サンプリングタスクとトレーニングタスク間のデータ変換とコピーを削減することを考慮する必要があります。サンプリングとトレーニングは典型的な生産者と消費者の関係を持っているため、異なるタスクに分割し、バッファを介して接続することで、生産者と消費者の処理能力のバランスを取り、データキャッシュとして機能し、分離を実現できます。生産者と消費者のアプローチに基づいて、GLT分散トレーニングには、サンプリングプロセスとトレーニングプロセスの2つの基本的なプロセスがあります。

サンプリングプロセス：分散近傍サンプリングと特徴量収集を担当します。サンプリング結果はサンプリングメッセージチャネルに送信され、その後の学習タスクに使用されます。

トレーニングプロセス：PyTorch DDPの分散トレーニングプロセスに対応します。通常、各トレーニングプロセスはトレーニングに1つのGPUを占有します。

これらのプロセスは、複数のマシンに柔軟に分散できます。分散プロセスデプロイメントをより適切に管理するために、GLTの分散トレーニングでは、ワーカーモードとサーバークライアントモードという2つのリファレンスデプロイメントモードを提供しています。

ワーカーモードでは、データの分割後、各マシンがシャードを保持し、サンプリングプロセスとトレーニングプロセスがこれらのマシンにまとめて展開されます。各トレーニングプロセスは複数のサンプリング子プロセスを生成でき、サンプリングプロセスは共有メモリメッセージチャネルを介してサンプリング結果をトレーニングプロセスに渡します。サンプリングプロセスでは、マルチプロセスサンプリングを使用でき、各分散サンプリング演算子はPythonコルーチンを使用して並行実行し、結果をメッセージチャネルに配置します。メッセージチャネルからトレーニングプロセスのGPUへのデータのコピーにかかる時間を短縮するために、メッセージチャネルをピンメモリに配置することもできます。

サーバー・クライアント・モードでは、クラスターにはサーバーノードとクライアントノードという2種類のマシンノードが含まれます。サンプリングプロセスはサーバーノードに展開され、トレーニングプロセスはすべてのクライアントノードに分散されます。サンプリングプロセスによって生成されたサンプル結果は、RPCを使用して実装されたリモートメッセージチャネルを介して、現在のトレーニングプロセスに送信されます。サーバー・クライアント・モードでは、サンプリングとトレーニングのワークロードを異なるマシンに配置することで、リソースの分離が可能です。

全体的なアーキテクチャ

LT の主な目的は、GPU/NVLink/RDMA などのハードウェア リソースと GNN モデルの特性を最大限に活用して、単一マシンおよび分散環境でのエンドツーエンドの GNN トレーニングを高速化することです。

ストレージ：GPUトレーニングのシナリオでは、グラフサンプリングとCPU-GPU間のデータ転送が大きなパフォーマンスボトルネックになりやすいです。グラフサンプリングと特徴量検索を高速化するために、GLTはCPUとGPUのメモリ管理を統合するUnifiedTensorを実装しています。特徴量収集によるCPU-GPU間のデータ転送オーバーヘッドを削減するため、GLTはホットな特徴量をGPUメモリにキャッシュし、残りの特徴量データにはUVA経由でアクセスすることをサポートしています。また、高速NVLinkを活用することで、GPU間のGPUキャッシュ容量を拡張しています。

グラフ操作演算子: ストレージ上で、GLT は、近傍サンプリング、ネガティブ サンプリング、特徴検索、サブグラフ サンプリングなど、CPU と GPU の両方をサポートするグラフ操作演算子を実装します。

分散オペレーター: 分散トレーニングでは、リモート データ アクセスによってモデルのトレーニング プロセスがブロックされるのを防ぐために、GLT は PyTorch RPC の上に効率的な RPC フレームワークをカプセル化し、マルチプロセスの並列および非同期の同時グラフ サンプリングと機能検索操作を使用してネットワークの遅延を隠し、エンドツーエンドのトレーニング スループットを向上させます。

インターフェース：PyGユーザーの学習曲線を短縮するため、GLTの上位API（ Sampler 、 Dataset 、 Loaderなど）はインターフェース面でPyGと互換性があります。そのため、PyGユーザーはわずかなコードの変更だけでGLTのアクセラレーション機能を最大限に活用できます。

モデル：GLTはPyGと互換性があるため、ほぼすべてのPyGモデルをベースモデルとして使用できます。さらに、豊富な分散学習サンプルも提供しています。

システムパフォーマンス

A100 GPUを搭載したマシンで、単一マシンのスケーラビリティテストを実施しました。テスト環境は、CUDA 11.4、PyTorch 1.12、GLT 0.2.0rc2です。下図は、近傍サンプリングと特徴量検索の合計スループットを示しています。GLTは線形スケーラビリティを備えていることがわかります（NVLinkにより、複数のGPUのキャッシュ容量が大きくなり、超線形の高速化が実現されます）。

さらに、分散サンプリングと特徴量検索のスケーラビリティを複数マシン間でテストしました。下の図は、各マシンに2基のA100 GPUを搭載した場合、2台マシンと4台マシンを1台マシンで実行した場合のスループット向上率を示しています。テストはCUDA 11.4、PyTorch 1.12、GLT 0.2.0rc2を使用して実施しました。2台マシンから4台マシンまで、ほぼ線形のスケーラビリティが見られることがわかります。

最後に、分散エンドツーエンド（e2e）モデルのパフォーマンスをテストしました。A100カードを2枚ずつ搭載したマシン2台で、DGL（DGLバージョン0.9.1、GLTバージョン0.2.0）との予備的な比較を行いました。

結論

この記事では、PyTorchをベースとしたGPUアクセラレーション対応分散GNNフレームワーク、GraphLearn-for-PyTorch (GLT) を紹介します。GLTは、分散GPU学習向けに最適化されたアクセラレーション機能を提供し、グラフサンプリングや特徴量検索処理にGPUなどのハードウェアリソースを最大限に活用し、線形スケーラビリティを特徴としています。上位層インターフェースはPyGと互換性があり、既存のPyGモデルを容易に高速化したり、既存のモデルを分散バージョンに変換したりできます。GLTはオープンソースであり、サンプルコードはPyGとGraphScopeで入手できます。私たちは今後もGLTの開発と最適化を続けていきますので、皆様のご利用と貢献をお待ちしております。

[1] P3: 大規模分散型ディープグラフ学習

[2] Quiver: ワークロード認識機能を備えた低レイテンシ、高スループットのGNNサービスのためのGPUサポート

無料のインタラクティブ モデリング パッケージ (PAI-DSW)、モデル トレーニング パッケージ (PAI-DLC 5000CU*H コンピューティング リソース)、および 500 元相当の無料のオンライン モデル サービス パッケージ (PAI-EAS) を入手できます。

https://free.aliyun.com/?product=9602825&tlog=0614GLTitpub