大規模ミニバッチGNN学習のための、通信不要なサンプリングと4Dハイブリッド並列性

要旨: グラフニューラルネットワーク（GNN）は、さまざまな実世界のシナリオから派生したグラフデータセットに対する学習で広く用いられています。非常に大規模なグラフから学習するには分散学習が必要であり、サンプリングを用いたミニバッチ処理は、GNN学習を並列化するための代表的な手法です。既存の分散ミニバッチ手法は、高コストなサンプリング手法や、データ並列性を用いる際のスケーリングの制限により、重要な性能ボトルネックを抱えています。本研究では、通信を必要としない分散サンプリング、3D並列行列乗算（PMM）、およびデータ並列性を組み合わせた、スケーラブルなミニバッチGNN学習のための4D並列フレームワークScaleGNNを提案します。ScaleGNNは統一的な頂点サンプリングアルゴリズムを導入し、各プロセス（GPUデバイス）が、プロセス間の通信なしで自分のローカルミニバッチ、すなわちサブグラフ分割を構築できるようにします。3D PMMにより、従来のデータ並列性（バニラ）よりも大幅に低い通信オーバーヘッドで、より多くのGPU数へミニバッチ学習をスケールさせることができます。さらに、サンプリングと学習のオーバーラップを可能にする最適化、低精度でデータを送信することによる通信オーバーヘッドの削減、カーネル融合、通信・計算のオーバーラップなどの追加最適化も示します。5つのグラフデータセットでScaleGNNを評価し、Perlmutterでは最大2048 GPUまで、Frontierでは2048 GCDsまで、Tuolumneでは最大1024 GPUまで強いスケーリングが可能であることを実証します。PerlmutterにおいてScaleGNNは、ogbn-productsでSOTA（最先端）ベースラインに対してエンドツーエンド学習速度を3.5倍向上させます。