[論文レビュー] Accelerating Training and Inference of Graph Neural Networks with Fast Sampling and Pipelining

Tim Kaler, Nickolas Stathas+5 more|arXiv (Cornell University)|Oct 16, 2021

Advanced Graph Neural NetworksComputer Science参考文献 35被引用数 21

ひとこと要約

本論文では、近傍サンプリングの最適化、共有メモリ並列バッチ準備、GPU計算とデータ転送のパイプライン化により、GNNの学習および推論を高速化するSALIENTというシステムを提案する。単一GPUでは標準のPyTorch-Geometricに対して3倍の高速化を達成し、16GPUでは8倍の並列高速化を実現した。ogbn-papers100Mで3層のGraphSAGEモデルを2.0秒/エポックで学習し、テスト精度は64.58%を達成した。

ABSTRACT

Improving the training and inference performance of graph neural networks (GNNs) is faced with a challenge uncommon in general neural networks: creating mini-batches requires a lot of computation and data movement due to the exponential growth of multi-hop graph neighborhoods along network layers. Such a unique challenge gives rise to a diverse set of system design choices. We argue in favor of performing mini-batch training with neighborhood sampling in a distributed multi-GPU environment, under which we identify major performance bottlenecks hitherto under-explored by developers: mini-batch preparation and transfer. We present a sequence of improvements to mitigate these bottlenecks, including a performance-engineered neighborhood sampler, a shared-memory parallelization strategy, and the pipelining of batch transfer with GPU computation. We also conduct an empirical analysis that supports the use of sampling for inference, showing that test accuracies are not materially compromised. Such an observation unifies training and inference, simplifying model implementation. We report comprehensive experimental results with several benchmark data sets and GNN architectures, including a demonstration that, for the ogbn-papers100M data set, our system SALIENT achieves a speedup of 3x over a standard PyTorch-Geometric implementation with a single GPU and a further 8x parallel speedup with 16 GPUs. Therein, training a 3-layer GraphSAGE model with sampling fanout (15, 10, 5) takes 2.0 seconds per epoch and inference with fanout (20, 20, 20) takes 2.4 seconds, attaining test accuracy 64.58%.

研究の動機と目的

• ミニバッチ準備およびデータ転送のコストが学習および推論の性能ボトルネックとなる、特に分散マルチGPU環境下での問題を解決すること。
• PyTorch-Geometric や DGL などの既存のGNNフレームワークで未解消のボトルネック、特にバッチ処理中のCPU利用率およびメモリ帯域幅の問題を特定・解決すること。
• 近傍サンプリングを用いて学習と推論を統合し、精度にほとんど影響を与えないことを実証的に検証することで、モデル実装の簡素化を図ること。
• データ転送と計算のパイプライン化により、GPUの利用効率を高め、エンドツーエンドの学習時間が唯一の遅延要因に制限されることを保証すること。
• PyTorchおよびPyGと互換性があり、他のフレームワークおよび大規模分散環境へ容易に拡張可能な汎用システムを構築すること。

提案手法

• CPUアーキテクチャに最適化されたパフォーマンス指向の近傍サンプラーの設計により、多ホップ近傍拡張におけるオーバーヘッドを低減すること。
• CPUおよびメモリ帯域幅のボトルネックを克服するため、バッチ準備における共有メモリ並列化の実装。
• CPUとGPU間のパイプライン化されたデータ転送を導入し、データ移動とGNN計算を重ねることでGPU利用効率を向上させること。
• 複数のマシンおよびGPUに跨る分散学習のために、PyTorchのDDPモジュールを活用し、最小限のコード変更で水平スケーリングを可能にすること。
• アーキテクチャの変更を必要とせず、既存のGNNアーキテクチャおよびフレームワーク（例：GraphSAGE）を活用することで、後方互換性および導入の容易性を確保すること。
• 学習および推論に一貫してサンプリングを適用しても顕著な精度低下がないことを実証的に評価し、その妥当性を検証すること。

実験結果

リサーチクエスチョン

• RQ1 マルチGPU環境下で近傍サンプリングを用いたGNN学習および推論における主なパフォーマンスボトルネックは何か？
• RQ2 バッチ準備およびデータ転送をどのように最適化すれば、全体の学習スループットへの影響を最小限に抑えられるか？
• RQ3 近傍サンプリングは推論精度にどの程度影響を与えるか？また、学習と推論に一貫して使用しても顕著な精度低下がないか？
• RQ4 GPU計算とデータ転送をパイプライン化することで、GNNワークロードにおいてほぼ理想に近いGPU利用効率を達成できるか？
• RQ5 提案手法は複数のGPUおよびマシンにどのようにスケーリング可能か？標準実装と比較して、達成可能な高速化はどの程度か？

主な発見

• SALIENTは、ogbn-papers100Mで3層のGraphSAGEモデルを学習する際、単一GPUで標準のPyTorch-Geometricに対して3倍の高速化を達成した。
• 16GPUを用いる場合、SALIENTは8倍の並列高速化を達成し、同じモデルおよびデータセットの学習時間を1エポックあたり2.0秒にまで短縮した。
• ファンアウト(20, 20, 20)を用いた推論では、1バッチあたり2.4秒で実行され、テスト精度は64.58%を達成した。これは、サンプリングによる精度低下が最小限であることを示している。
• バッチ準備、データ転送、GPU計算の間でほぼ完全なオーバーラップが実現されており、エンドツーエンドの学習時間は、個々のコンponentの遅延の最大値にほぼ等しくなっている。
• 最適化されたサンプリングおよび共有メモリ並列化により、特にCPUおよびメモリ帯域幅のボトルネックが顕著に低減された。
• 提案された最適化は汎用的であり、他のGNNフレームワークおよび分散環境へも適用可能であり、GPUベースのスライシングやデータキャッシュなどの技術とさらに統合可能である可能性を有する。