Core Concepts
グラフニューラルネットワークの分散トレーニングにおいて、サンプリングステップの通信コストを削減する新しい手法を提案する。行列ベースのバルクサンプリングアプローチを用いて、複数のミニバッチをまとめてサンプリングすることで、通信コストを大幅に削減できる。
Abstract
本論文では、グラフニューラルネットワーク(GNN)のトレーニングにおける分散サンプリングの新しい手法を提案している。
GNNモデルは大規模であるため、分散ミニバッチトレーニングが必要となる。本論文の主要な貢献は、分散GNNトレーニングにおけるサンプリングステップの通信コストを削減する新しい手法である。
提案手法では、行列ベースのバルクサンプリングアプローチを用いて、複数のミニバッチをまとめてサンプリングする。入力グラフの構造がシングルデバイスに収まらない場合は、グラフを分散させ、通信を回避するスパース行列乗算アルゴリズムを使ってスケーリングする。入力グラフの構造(ただし埋め込みは除く)がシングルGPUのメモリに収まる場合は、(1)通信なしでサンプリングを行い、(2)ミニバッチのサンプリングのオーバーヘッドを軽減し、(3)行列構築を変えるだけで複数のサンプリングアルゴリズムを表現できる。
サンプリングの新手法に加えて、提案するパイプラインはマトリックスベースのバルクサンプリングアプローチを使ってエンドツーエンドのトレーニング結果を提供する。Open Graph Benchmark (OGB)の最大データセットで128GPUを使った実験結果を示し、提案パイプラインがQuiverに比べて2.5倍高速であることを示す。OGB以外のデータセットでは、128GPUで1エポックあたり8.46倍高速化を示す。さらに、グラフを分散させた場合や、ノードワイズとレイヤーワイズのサンプリングアルゴリズムの両方でスケーリングすることを示す。
Stats
提案手法は、Quiverに比べて、Products データセットで16GPUで2.5倍高速、Papers データセットで64GPUで3.4倍高速、Protein データセットで128GPUで8.5倍高速である。
OGB以外のデータセットでは、128GPUで1エポックあたり8.46倍高速化された。