Sequential Aggregation and Rematerialization: Distributed Full-batch Training of Graph Neural Networks on Large Graphs

Abstract

We present the Sequential Aggregation and Rematerialization (SAR) scheme for distributed full-batch training of Graph Neural Networks (GNNs) on large graphs. Large-scale training of GNNs has recently been dominated by sampling-based methods and methods based on non-learnable message passing. SAR on the other hand is a distributed technique that can train any GNN type directly on an entire large graph. The key innovation in SAR is the distributed sequential rematerialization scheme which sequentially re-constructs then frees pieces of the prohibitively large GNN computational graph during the backward pass. This results in excellent memory scaling behavior where the memory consumption per worker goes down linearly with the number of workers, even for densely connected graphs. Using SAR, we report the largest applications of full-batch GNN training to-date, and demonstrate large memory savings as the number of workers increases. We also present a general technique based on kernel fusion and attention-matrix rematerialization to optimize both the runtime and memory efficiency of attention-based models. We show that, coupled with SAR, our optimized attention kernels lead to significant speedups and memory savings in attention-based GNNs.We made the SAR GNN training library publicy available: \url{https://github.com/IntelLabs/SAR}.

References in corpus (16)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Neural Message Passing for Quantum Chemistry
• Simplifying Graph Convolutional Networks
• Representation Learning on Graphs with Jumping Knowledge Networks
• Open Graph Benchmark: Datasets for Machine Learning on Graphs
• FastGCN: Fast Learning with Graph Convolutional Networks via Importance Sampling
• GraphSAINT: Graph Sampling Based Inductive Learning Method
• Link Prediction Based on Graph Neural Networks
• DeeperGCN: All You Need to Train Deeper GCNs
• Graph Matching Networks for Learning the Similarity of Graph Structured Objects
• The Reversible Residual Network: Backpropagation Without Storing Activations
• SIGN: Scalable Inception Graph Neural Networks
• Combining Label Propagation and Simple Models Out-performs Graph Neural Networks
• Masked Label Prediction: Unified Message Passing Model for Semi-Supervised Classification
• Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization
• Scalable and Adaptive Graph Neural Networks with Self-Label-Enhanced training