[論文レビュー] Sequential Aggregation and Rematerialization: Distributed Full-batch Training of Graph Neural Networks on Large Graphs

Hesham Mostafa|arXiv (Cornell University)|Nov 11, 2021

Advanced Graph Neural Networks被引用数 4

ひとこと要約

本稿では、順次再構築と再物質化（SAR）を導入し、バックプロパゲーション中に計算グラフの部品を順次再構築・解放することで、メモリ使用量を著しく削減する分散型フルバッチ学習手法を提案する。SARは、128ワーカーで運用する際、最大50％のメモリ削減を達成し、最適化されたアテンションカーネルを用いることで顕著な高速化を実現する。これにより、ogbn-papers100M や ogbn-products といった大規模グラフの学習が可能になる。

ABSTRACT

We present the Sequential Aggregation and Rematerialization (SAR) scheme for distributed full-batch training of Graph Neural Networks (GNNs) on large graphs. Large-scale training of GNNs has recently been dominated by sampling-based methods and methods based on non-learnable message passing. SAR on the other hand is a distributed technique that can train any GNN type directly on an entire large graph. The key innovation in SAR is the distributed sequential rematerialization scheme which sequentially re-constructs then frees pieces of the prohibitively large GNN computational graph during the backward pass. This results in excellent memory scaling behavior where the memory consumption per worker goes down linearly with the number of workers, even for densely connected graphs. Using SAR, we report the largest applications of full-batch GNN training to-date, and demonstrate large memory savings as the number of workers increases. We also present a general technique based on kernel fusion and attention-matrix rematerialization to optimize both the runtime and memory efficiency of attention-based models. We show that, coupled with SAR, our optimized attention kernels lead to significant speedups and memory savings in attention-based GNNs.We made the SAR GNN training library publicy available: \\url{https://github.com/IntelLabs/SAR}.

研究の動機と目的

• 大規模かつ密に接続されたグラフにおいて、計算グラフが保存するのに現実的でないほど大きくなるため、フルバッチGNN学習におけるメモリボトルネックを解消すること。
• サンプリングや非学習的メッセージパッシングに依存せずに、任意のGNNアーキテクチャをスケーラブルかつ分散型でフルバッチ学習可能にすること。
• 特にアテンションベースのモデルにおいて、中間テンソルの重複した物質化を回避することで、分散学習における通信オーバーヘッドを最小限に抑えること。
• サンプリングベースや非学習的メッセージパッシングGNN手法の評価のための実用的でメモリ効率の良いベースラインを提供すること。
• GATのようなアテンションベースのGNNのパフォーマンスを最適化するために、アテンション係数の保存を排除し、オンザフライで計算すること。

提案手法

• SARは、入力グラフをNワーカーに分散して分割し、各ワーカーが割り当てられた部分グラフのみを処理する分散ドメイン並列学習方式を採用する。
• フォワードパス中に計算グラフ全体を物質化するのではなく、SARはバックプロパゲーションまでグラフ構築を延期し、部分的に再構築する。
• バックプロパゲーション中、SARは計算グラフ部品を順次再物質化し、解放することで、任意の時点で1ワーカーあたり最大2つのパーティションしか保持しないように保証する。
• この方法により、ワーカーあたりのメモリスケーリングが O(2/N) に抑えられ、グラフの密度に依存せず、より多くのワーカーを追加することで任意に大きなグラフの学習が可能になる。
• アテンションベースのモデルでは、SARがカーネル統合とオンザフライでのアテンション係数計算を統合し、大きなアテンション行列の保存を回避する。
• 空間並列CNNなど、出力が跨ワーカー入力に依存するような任意のドメイン並列学習設定にも一般化可能である。

実験結果

リサーチクエスチョン

• RQ1 サンプリングや非学習的メッセージパッシングに依存せずに、大規模グラフにおけるフルバッチGNN学習をスケーリング可能か？
• RQ2 分散GNN学習におけるメモリ使用量をどのように低減すれば、メモリに収まらないほど大きなグラフの学習が可能になるか？
• RQ3 分散環境下でバックプロパゲーション中に計算グラフを再物質化する際の通信およびメモリオーバーヘッドはどの程度か？
• RQ4 GATモデルのような高価な中間テンソル（例：アテンション係数）を再物質化しないために、再物質化を最適化できるか？
• RQ5 大規模ベンチマークにおいて、SARはサンプリングベースや非学習的メッセージパッシングGNN手法と比較して、メモリ効率および学習速度で優れているか？

主な発見

• ogbn-papers100M で128ワーカーで GraphSage モデルを学習する際、SARはワーカーあたりのピークメモリ使用量を最大50％削減し、2/N の線形的スケーリングを達成した。
• ogbn-products では、最適化されたアテンションカーネルを用いることで、DGLのGAT実装に比べ2.5倍の高速化を達成した。これは、メモリ圧迫の低減とオンザフライ係数計算によるものである。
• GraphSage では、通信量を低く抑えつつ、ドメイン並列学習と同等の実行時間で、128ワーカーでメモリ使用量を半分に削減した。
• SARと併用された最適化アテンションカーネル（FAK）は、アテンション係数の保存を回避することでメモリ使用量を削減し、バックプロパゲーションのパフォーマンスに悪影響を与えずにフォワードパスの速度を向上させた。
• SARは、ogbn-papers100M および ogbn-products における最大の報告済みフルバッチGNN学習を実現し、1億個以上のノードを持つグラフでも実行可能であることを示した。
• 多くのGNNバリアントにおいて通信を回避するため、メモリの節約が実際の実行時間のオーバーヘッドなしに達成可能である。