[論文レビュー] Boosting Distributed Full-graph GNN Training with Asynchronous One-bit Communication

Meng Zhang, Qinghao Hu+3 more|arXiv (Cornell University)|Mar 2, 2023

Advanced Graph Neural Networks被引用数 4

ひとこと要約

本稿では、全グラフGNN学習におけるGPU間通信のオーバーヘッドを著しく削減するために、非同期1ビット通信を用いる分散GNN学習フレームワークSylvieを提案する。低ビット量子化および誤差補償を用いた復元により、Sylvieはベースライン手法と比較して最大28.1倍の高速化を達成しながら、モデルの精度を維持する。

ABSTRACT

Training Graph Neural Networks (GNNs) on large graphs is challenging due to the conflict between the high memory demand and limited GPU memory. Recently, distributed full-graph GNN training has been widely adopted to tackle this problem. However, the substantial inter-GPU communication overhead can cause severe throughput degradation. Existing communication compression techniques mainly focus on traditional DNN training, whose bottleneck lies in synchronizing gradients and parameters. We find they do not work well in distributed GNN training as the barrier is the layer-wise communication of features during the forward pass & feature gradients during the backward pass. To this end, we propose an efficient distributed GNN training framework Sylvie, which employs one-bit quantization technique in GNNs and further pipelines the curtailed communication with computation to enormously shrink the overhead while maintaining the model quality. In detail, Sylvie provides a lightweight Low-bit Module to quantize the sent data and dequantize the received data back to full precision values in each layer. Additionally, we propose a Bounded Staleness Adaptor to control the introduced staleness to achieve further performance enhancement. We conduct theoretical convergence analysis and extensive experiments on various models & datasets to demonstrate Sylvie can considerably boost the training throughput by up to 28.1x.

研究の動機と目的

• 分散型全グラフGNN学習における高いGPU間通信オーバーヘッドを低減し、スケーラビリティとスループットを向上させること。
• 従来の圧縮技術には限界があり、層ごとの特徴量と勾配の通信に特化したDNN向けに最適化されているが、GNNには効果的でないことを克服すること。
• 低ビット量子化と非同期パイプライン実行を用いて、通信コストを著しく削減しながらも高いモデル精度を維持すること。
• 収束性やパフォーマンスを損なわず、多様なGNNモデルとデータセットに一般化可能なシステムを設計すること。

提案手法

• 送信時に特徴量と勾配を1ビットに量子化し、受信側で誤差補償付き再構成によりフル精度に復元する低ビットモジュールを導入する。
• 通信と計算を重ねて実行する非同期パイプライン（Sylvie-A）を実装し、レイテンシを隠蔽しリソース利用効率を向上させる。
• 非同期学習における古くなった特徴量と勾配の影響を制御するためのバウンデッドスタリネスアダプタを提案し、収束安定性を向上させる。
• グラフをGPU間で分割し、通信が重要なノードを同定することで、重複するデータ転送を最小限に抑えるグラフエンジンを設計する。
• 復元されたデータを用いて前方伝搬と逆伝搬を実行するトレーナーモジュールを採用し、量子化の影響を受けてもモデルの一貫性を保証する。
• 通信に必要なテンソルにのみ量子化を適用し、モデルコア部ではフル精度の計算を維持する。

実験結果

リサーチクエスチョン

• RQ1 1ビット量子化は、モデル精度を劣化させることなく、分散GNN学習における通信オーバーヘッドを効果的に削減できるか？
• RQ2 低ビット通信と非同期パイプライン実行は、GNNの学習スループットと収束性にどのように影響を与えるか？
• RQ3 特徴量と勾配のスタリネスは、モデルの収束にどの程度影響を及ぼし、性能を維持するためにどのように制御できるか？
• RQ4 提案されたフレームワークは、異なるGNNアーキテクチャと実世界のグラフデータセットに一般化可能か？
• RQ5 低ビットGNN学習において、通信削減、計算オーバーヘッド、メモリフットプリントのトレードオフはどのようなものか？

主な発見

• Redditデータセットにおいて、SylvieはFP32ベースラインと比較して通信オーバーヘッドを89.8％削減し、1エポックあたりの学習時間を84.2％短縮した。
• 1ビット量子化を用いることで、2サーバー構成においてSylvieは最大28.1倍の学習スループット向上を達成した（図1参照）。
• すべての埋め込みと勾配を1ビットに量子化すると、深刻な精度低下（例：Ogbn-productsでは70.6％ vs. Sylvie-Sの78.86％）が生じ、選択的量子化が不可欠であることを示した。
• 2サーバー構成において、低ビットモジュールの総学習時間に占める寄与はわずか5.9％にとどまり、計算オーバーヘッドが無視できるレベルであることがわかった。
• バウンデッドスタリネスアダプタは、非同期更新が性能を劣化させる状況においても収束性を向上させた。
• 複数のモデル（GraphSAGE, GCN, GAT）とデータセット（Reddit, Yelp, Ogbn-products）において、Sylvieはフル精度学習と比較して最小限の精度損失で高い精度を維持した。