Tango: rethinking quantization for graph neural network training on GPUs

TL;DR

Tango presents a new approach to quantized GNN training on GPUs that maintains accuracy and accelerates training by introducing efficient quantization rules, optimized primitives, and system integration, outperforming existing methods.

Contribution

The paper introduces novel quantization rules, optimized primitives, and system integration techniques to improve accuracy and speed in quantized GNN training on GPUs.

Findings

Tango achieves faster training times than state-of-the-art methods.

Tango maintains high accuracy in quantized GNN training.

System integration with DGL enhances practical applicability.

Abstract

Graph Neural Networks (GNNs) are becoming increasingly popular due to their superior performance in critical graph-related tasks. While quantization is widely used to accelerate GNN computation, quantized training faces unprecedented challenges. Current quantized GNN training systems often have longer training times than their full-precision counterparts for two reasons: (i) addressing the accuracy challenge leads to excessive overhead, and (ii) the optimization potential exposed by quantization is not adequately leveraged. This paper introduces Tango which re-thinks quantization challenges and opportunities for graph neural network training on GPUs with three contributions: Firstly, we introduce efficient rules to maintain accuracy during quantized GNN training. Secondly, we design and implement quantization-aware primitives and inter-primitive optimizations that can speed up GNN training. Finally, we integrate Tango with the popular Deep Graph Library (DGL) system and demonstrate its superior performance over state-of-the-art approaches on various GNN models and datasets.