# Distributed Matrix-Based Sampling for Graph Neural Network Training

arXiv: https://arxiv.org/abs/2311.02909 · HTML: https://mlsanity.com/paper/2311.02909/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 38 references that are in the map.
- First posted on arXiv 6 Nov 2023.

## Builds on (most cited first, 19 of 19)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [A Comprehensive Survey on Graph Neural Networks](https://mlsanity.com/paper/1901.00596/) (2019)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [GraphSAINT: Graph Sampling Based Inductive Learning Method](https://mlsanity.com/paper/1907.04931/) (2019)
- [FastGCN: Fast Learning with Graph Convolutional Networks via Importance Sampling](https://mlsanity.com/paper/1801.10247/) (2018)
- [Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks](https://mlsanity.com/paper/1905.07953/) (2019)
- [Layer-Dependent Importance Sampling for Training Deep and Large Graph Convolutional Networks](https://mlsanity.com/paper/1911.07323/) (2019)
- [AliGraph: A Comprehensive Graph Neural Network Platform](https://mlsanity.com/paper/1902.08730/) (2019)
- [DistDGL: Distributed Graph Neural Network Training for Billion-Scale Graphs](https://mlsanity.com/paper/2010.05337/) (2020)
- [Reducing Communication in Graph Neural Network Training](https://mlsanity.com/paper/2005.03300/) (2020)
- [PipeGCN: Efficient Full-Graph Training of Graph Convolutional Networks with Pipelined Feature Communication](https://mlsanity.com/paper/2203.10428/) (2022)
- [BNS-GCN: Efficient Full-Graph Training of Graph Convolutional Networks with Partition-Parallelism and Random Boundary Node Sampling](https://mlsanity.com/paper/2203.10983/) (2022)
- [Accelerating Graph Sampling for Graph Machine Learning using GPUs](https://mlsanity.com/paper/2009.06693/) (2020)
- [Integrated Model, Batch and Domain Parallelism in Training Neural Networks](https://mlsanity.com/paper/1712.04432/) (2017)
- [Communication-Efficient Graph Neural Networks with Probabilistic Neighborhood Expansion Analysis and Caching](https://mlsanity.com/paper/2305.03152/) (2023)
- [C-SAW: A Framework for Graph Sampling and Random Walk on GPUs](https://mlsanity.com/paper/2009.09103/) (2020)
- [GraphBLAST: A High-Performance Linear Algebra-based Graph Framework on the GPU](https://mlsanity.com/paper/1908.01407/) (2019)

## Built on by (most cited first, 6 of 6)

- [MassiveGNN: Efficient Training via Prefetching for Massively Connected Distributed Graphs](https://mlsanity.com/paper/2410.22697/) (2024)
- [Scaling Graph Neural Networks for Particle Track Reconstruction](https://mlsanity.com/paper/2504.04670/) (2025)
- [GPU-centric Communication Schemes for HPC and ML Applications](https://mlsanity.com/paper/2503.24230/) (2025)
- [The Ubiquitous Sparse Matrix-Matrix Products](https://mlsanity.com/paper/2508.04077/) (2025)
- [Communication-Avoiding Linear Algebraic Kernel K-Means on GPUs](https://mlsanity.com/paper/2601.17136/) (2026)
- [Communication-free Sampling and 4D Hybrid Parallelism for Scalable Mini-batch GNN Training](https://mlsanity.com/paper/2604.02651/) (2026)

## Abstract (arXiv)

Graph Neural Networks (GNNs) offer a compact and computationally efficient way to learn embeddings and classifications on graph data. GNN models are frequently large, making distributed minibatch training necessary. The primary contribution of this paper is new methods for reducing communication in the sampling step for distributed GNN training. Here, we propose a matrix-based bulk sampling approach that expresses sampling as a sparse matrix multiplication (SpGEMM) and samples multiple minibatches at once. When the input graph topology does not fit on a single device, our method distributes the graph and use communication-avoiding SpGEMM algorithms to scale GNN minibatch sampling, enabling GNN training on much larger graphs than those that can fit into a single device memory. When the input graph topology (but not the embeddings) fits in the memory of one GPU, our approach (1) performs sampling without communication, (2) amortizes the overheads of sampling a minibatch, and (3) can represent multiple sampling algorithms by simply using different matrix constructions. In addition to new methods for sampling, we introduce a pipeline that uses our matrix-based bulk sampling approach to provide end-to-end training results. We provide experimental results on the largest Open Graph Benchmark (OGB) datasets on 128 GPUs, and show that our pipeline is 2.5× faster than Quiver (a distributed extension to PyTorch-Geometric) on a 3-layer GraphSAGE network. On datasets outside of OGB, we show a 8.46× speedup on 128 GPUs in per-epoch time. Finally, we show scaling when the graph is distributed across GPUs and scaling for both node-wise and layer-wise sampling algorithms.