[論文レビュー] QGTC: Accelerating Quantized Graph Neural Networks via GPU Tensor Core

Yuke Wang, Boyuan Feng+1 more|arXiv (Cornell University)|Nov 18, 2021

Advanced Graph Neural Networks被引用数 4

ひとこと要約

QGTC は、GPU Tensor Core を活用した新規フレームワークであり、量子化グラフニューラルネットワーク（QGNN）における高性能で任意のビット幅の推論を実現する。ビット分解算術、3次元スタックドビット圧縮、ゼロタイルジャンプ、非ゼロタイル再利用を、TC最適化されたCUDAカーネルに統合することで、多様なGNNワークロードにおいて、最先端のDGLを平均2.7倍高速化し、スパースで低ビットのGNN計算におけるGPU Tensor Coreの有効な活用を示している。

ABSTRACT

Over the most recent years, quantized graph neural network (QGNN) attracts lots of research and industry attention due to its high robustness and low computation and memory overhead. Unfortunately, the performance gains of QGNN have never been realized on modern GPU platforms. To this end, we propose the first Tensor Core (TC) based computing framework, QGTC, to support any-bitwidth computation for QGNNs on GPUs. We introduce a novel quantized low-bit arithmetic design based on the low-bit data representation and bit-decomposed computation. We craft a novel TC-tailored CUDA kernel design by incorporating 3D-stacked bit compression, zero-tile jumping, and non-zero tile reuse technique to improve the performance systematically. We incorporate an effective bandwidth-optimized subgraph packing strategy to maximize the transferring efficiency between CPU host and GPU device. We integrate QGTC with Pytorch for better programmability and extensibility. Extensive experiments demonstrate that QGTC achieves on average 2.7x speedup compared with the state-of-the-art Deep Graph Library framework across diverse settings.

研究の動機と目的

• 低メモリおよび計算オーバーヘッドであるにもかかわらず、現代のGPU上で量子化GNN（QGNN）をデプロイする際のパフォーマンスギャップを解消すること。
• CUDAコアに依存する既存のGPUフレームワークでは、低精度GNNのビットレベル演算を効率的に活用できないという制限を克服すること。
• GPU Tensor Cores上で任意のビット幅（例：1-, 2-, 4ビット）の計算を可能にすること。現在のGNNワークロードでは、これらが未活用のままである。
• PyTorchなどのハイレベルディープラーニングフレームワークとの互換性を保ちつつ、低ビットGNN計算とを接続するソフトウェアスタックを設計すること。
• スパースGNNワークロードのための帯域幅最適化された部分グラフパッケージングとTC特化カーネル最適化により、GPUの利用度を最大化すること。

提案手法

• ネイティブに1ビットおよび4ビットのTensor Core演算をサポートするものに基づき、任意のビット幅GNN計算をエミュレートするビット分解算術設計を提案する。
• 複数の低ビット値を1つの32ビットワードに圧縮する3次元スタックドビット圧縮技術を実装し、メモリフットプリントの削減とデータ再利用の向上を図る。
• ゼロ値の隣接行列タイルの計算をスキップするゼロタイルジャンプを備えた、TC最適化されたCUDAカーネルを設計する。これにより、不要なメモリアクセスが最小限に抑えられる。
• 複数のGEMM演算にわたり、事前に読み込んだ非ゼロの隣接行列タイルを再利用する非ゼロタイル再利用を導入し、グローバルメモリ帯域幅の圧力を軽減する。
• 同様のサイズとアクセスパターンを持つ部分グラフを統合することで、CPUホストとGPUデバイス間のデータ転送効率を最大化する帯域幅最適化部分グラフパッケージング戦略を開発する。
• PyTorchとの統合により、プログラマビリティ、拡張性、および既存のディープラーニングワークフローとの互換性を確保する。

実験結果

リサーチクエスチョン

• RQ1 GPU Tensor Cores は、従来のCUDAコアベースのフレームワークを凌駆えるものとして、量子化GNN推論を効果的に加速できるか？
• RQ2 1ビットおよび4ビットの演算のみをネイティブにサポートするハードウェア上で、2ビットや4ビットのような任意のビット幅GNN計算を効率的にエミュレートする方法は何か？
• RQ3 スパースかつ不規則なGNN計算におけるTensor Coresのパフォーマンスボトルネックを緩和するために、最も効果的なカーネルレベル最適化は何か？
• RQ4 部分グラフパッケージングとメモリアクセス最適化は、GNN推論におけるCPUとGPU間のデータ転送効率をどの程度向上できるか？
• RQ5 ビットレベル圧縮、ゼロタイルスキップ、非ゼロタイル再利用の組み合わせが、QGNN推論全体のスループットおよびGPU利用度に及ぼす影響は何か？

主な発見

• QGTC は、多様なGNNモデルおよびデータセットにおいて、最先端のDGLフレームワークを平均2.7倍高速化する推論性能を達成した。
• QGTCの1ビットGEMMカーネルのパフォーマンスは、部分グラフサイズの増加に伴い良好にスケーリングされ、特に512～16,384ノードの範囲で、計算強度の上昇に伴いGPU SM利用度が向上する。
• 非ゼロタイル再利用は、特にノード埋め込みに高ビット幅（4～16ビット）が使用される大規模行列（例：N=8,192）において、グローバルメモリアクセスのオーバーヘッドを低減し、スループットの向上に顕著に寄与する。
• 3次元スタックドビット圧縮技術は、メモリフットプリントの削減と効率的なビットレベル並列処理を実現し、Tensor Coresにおける高い算術強度の向上に貢献している。
• 帯域幅最適化された部分グラフパッケージングにより、CPUとGPU間のデータ転送効率が向上し、アイドルタイムが削減され、全体のシステムスループットが向上した。
• フレームワークは、現代のGPU上で優れたスケーラビリティとリソース利用度を示しており、大規模な部分グラフサイズにおいて、SMがフルにアクティブ化されることで、ピーク性能に近いパフォーマンスを発揮した。