Accelerating Training and Inference of Graph Neural Networks with Fast Sampling and Pipelining

TL;DR

This paper introduces system-level improvements for graph neural network training and inference, including sampling, pipelining, and parallelization, achieving significant speedups on large-scale datasets without sacrificing accuracy.

Contribution

It proposes a set of techniques to accelerate GNN training and inference in multi-GPU environments, addressing bottlenecks in mini-batch preparation and data transfer.

Findings

Achieves 3x speedup over standard PyTorch-Geometric on single GPU

Attains 8x parallel speedup with 16 GPUs

Maintains comparable accuracy during inference with sampling

Abstract

Improving the training and inference performance of graph neural networks (GNNs) is faced with a challenge uncommon in general neural networks: creating mini-batches requires a lot of computation and data movement due to the exponential growth of multi-hop graph neighborhoods along network layers. Such a unique challenge gives rise to a diverse set of system design choices. We argue in favor of performing mini-batch training with neighborhood sampling in a distributed multi-GPU environment, under which we identify major performance bottlenecks hitherto under-explored by developers: mini-batch preparation and transfer. We present a sequence of improvements to mitigate these bottlenecks, including a performance-engineered neighborhood sampler, a shared-memory parallelization strategy, and the pipelining of batch transfer with GPU computation. We also conduct an empirical analysis that supports the use of sampling for inference, showing that test accuracies are not materially compromised. Such an observation unifies training and inference, simplifying model implementation. We report comprehensive experimental results with several benchmark data sets and GNN architectures, including a demonstration that, for the ogbn-papers100M data set, our system SALIENT achieves a speedup of 3x over a standard PyTorch-Geometric implementation with a single GPU and a further 8x parallel speedup with 16 GPUs. Therein, training a 3-layer GraphSAGE model with sampling fanout (15, 10, 5) takes 2.0 seconds per epoch and inference with fanout (20, 20, 20) takes 2.4 seconds, attaining test accuracy 64.58%.