[Paper Review] Accelerating Training and Inference of Graph Neural Networks with Fast Sampling and Pipelining

Tim Kaler, Nickolas Stathas+5 more|arXiv (Cornell University)|Oct 16, 2021

Advanced Graph Neural Networks35 references21 citations

TL;DR

This paper proposes SALIENT, a system that accelerates GNN training and inference by optimizing neighborhood sampling, enabling shared-memory parallel batch preparation, and pipelining data transfer with GPU computation. It achieves 3× speedup over standard PyTorch-Geometric on a single GPU and 8× parallel speedup on 16 GPUs for ogbn-papers100M, training a 3-layer GraphSAGE model in 2.0 seconds per epoch with 64.58% test accuracy.

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%.

Motivation & Objective

• To address the performance bottleneck in GNN training and inference caused by expensive mini-batch preparation and data transfer, especially in distributed multi-GPU settings.
• To identify and resolve under-explored bottlenecks in existing GNN frameworks like PyTorch-Geometric and DGL, particularly in CPU utilization and memory bandwidth during batch processing.
• To unify training and inference through neighborhood sampling by empirically validating that accuracy is minimally affected, simplifying model implementation.
• To enable high GPU utilization through pipelining of data transfer and computation, ensuring end-to-end training time is limited only by the slowest component.
• To build a general-purpose system compatible with PyTorch and PyG that can be easily extended to other frameworks and large-scale distributed environments.

Proposed method

• Designing a performance-engineered neighborhood sampler that optimizes for CPU architecture and reduces overhead in multi-hop neighborhood expansion.
• Implementing shared-memory parallelization for batch preparation to overcome CPU and memory bandwidth bottlenecks in existing frameworks.
• Introducing pipelined data transfer between CPU and GPU, overlapping data movement with GNN computation to improve GPU utilization.
• Using PyTorch's DDP module for distributed training across multiple machines and GPUs, enabling horizontal scaling with minimal code changes.
• Leveraging existing GNN architectures and frameworks (e.g., GraphSAGE) without requiring architectural changes, ensuring backward compatibility and ease of adoption.
• Empirically evaluating sampling impact on inference accuracy to validate the feasibility of using sampling uniformly across training and inference.

Experimental results

Research questions

• RQ1 What are the primary performance bottlenecks in GNN training and inference when using neighborhood sampling on multi-GPU systems?
• RQ2 How can batch preparation and data transfer be optimized to reduce their impact on overall training throughput?
• RQ3 To what extent does neighborhood sampling affect inference accuracy, and can it be used uniformly across training and inference without significant accuracy loss?
• RQ4 Can pipelining data transfer with GPU computation achieve near-ideal GPU utilization in GNN workloads?
• RQ5 How does the proposed system scale across multiple GPUs and machines, and what is the achievable speedup compared to standard implementations?

Key findings

• SALIENT achieves a 3× speedup over standard PyTorch-Geometric on a single GPU when training a 3-layer GraphSAGE model on ogbn-papers100M.
• With 16 GPUs, SALIENT achieves an 8× parallel speedup, reducing training time to 2.0 seconds per epoch for the same model and dataset.
• Inference with fanout (20, 20, 20) takes 2.4 seconds per batch and attains a test accuracy of 64.58%, demonstrating minimal accuracy degradation from sampling.
• The system achieves near-perfect overlap between batch preparation, data transfer, and GPU computation, making end-to-end training time nearly equal to the slowest individual component.
• The performance bottlenecks in batch preparation and transfer are significantly reduced through optimized sampling and shared-memory parallelization, especially on CPU and memory bandwidth.
• The proposed optimizations are general and can be applied to other GNN frameworks and distributed environments, with potential for further integration with techniques like GPU-based slicing or data caching.