# PyTorch-Direct: Enabling GPU Centric Data Access for Very Large Graph Neural Network Training with Irregular Accesses

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 14 of its 28 references that are in the map.
- First posted on arXiv 20 Jan 2021.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [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)
- [Graph Neural Networks: A Review of Methods and Applications](https://mlsanity.com/paper/1812.08434/) (2018)
- [Graph Convolutional Neural Networks for Web-Scale Recommender Systems](https://mlsanity.com/paper/1806.01973/) (2018)
- [GraphSAINT: Graph Sampling Based Inductive Learning Method](https://mlsanity.com/paper/1907.04931/) (2019)
- [Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks](https://mlsanity.com/paper/1905.07953/) (2019)
- [Graph Neural Networks in TensorFlow and Keras with Spektral](https://mlsanity.com/paper/2006.12138/) (2020)
- [EMOGI: Efficient Memory-access for Out-of-memory Graph-traversal In GPUs](https://mlsanity.com/paper/2006.06890/) (2020)

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

- [IGB: Addressing The Gaps In Labeling, Features, Heterogeneity, and Size of Public Graph Datasets for Deep Learning Research](https://mlsanity.com/paper/2302.13522/) (2023)
- [MariusGNN: Resource-Efficient Out-of-Core Training of Graph Neural Networks](https://mlsanity.com/paper/2202.02365/) (2022)
- [TC-GNN: Bridging Sparse GNN Computation and Dense Tensor Cores on GPUs](https://mlsanity.com/paper/2112.02052/) (2021)
- [MGG: Accelerating Graph Neural Networks with Fine-grained intra-kernel Communication-Computation Pipelining on Multi-GPU Platforms](https://mlsanity.com/paper/2209.06800/) (2022)
- [Quiver: Supporting GPUs for Low-Latency, High-Throughput GNN Serving with Workload Awareness](https://mlsanity.com/paper/2305.10863/) (2023)
- [Accelerating Sampling and Aggregation Operations in GNN Frameworks with GPU Initiated Direct Storage Accesses](https://mlsanity.com/paper/2306.16384/) (2023)
- [FreshGNN: Reducing Memory Access via Stable Historical Embeddings for Graph Neural Network Training](https://mlsanity.com/paper/2301.07482/) (2023)
- [DGI: Easy and Efficient Inference for GNNs](https://mlsanity.com/paper/2211.15082/) (2022)
- [Russo-Ukrainian War: Prediction and explanation of Twitter suspension](https://mlsanity.com/paper/2306.03502/) (2023)
- [LSM-GNN: Large-scale Storage-based Multi-GPU GNN Training by Optimizing Data Transfer Scheme](https://mlsanity.com/paper/2407.15264/) (2024)
- [Armada: Memory-Efficient Distributed Training of Large-Scale Graph Neural Networks](https://mlsanity.com/paper/2502.17846/) (2025)
- [MQ-GNN: A Multi-Queue Pipelined Architecture for Scalable and Efficient GNN Training](https://mlsanity.com/paper/2601.04707/) (2026)

## Abstract (arXiv)

With the increasing adoption of graph neural networks (GNNs) in the machine learning community, GPUs have become an essential tool to accelerate GNN training. However, training GNNs on very large graphs that do not fit in GPU memory is still a challenging task. Unlike conventional neural networks, mini-batching input samples in GNNs requires complicated tasks such as traversing neighboring nodes and gathering their feature values. While this process accounts for a significant portion of the training time, we find existing GNN implementations using popular deep neural network (DNN) libraries such as PyTorch are limited to a CPU-centric approach for the entire data preparation step. This "all-in-CPU" approach has negative impact on the overall GNN training performance as it over-utilizes CPU resources and hinders GPU acceleration of GNN training. To overcome such limitations, we introduce PyTorch-Direct, which enables a GPU-centric data accessing paradigm for GNN training. In PyTorch-Direct, GPUs are capable of efficiently accessing complicated data structures in host memory directly without CPU intervention. Our microbenchmark and end-to-end GNN training results show that PyTorch-Direct reduces data transfer time by 47.1% on average and speeds up GNN training by up to 1.6x. Furthermore, by reducing CPU utilization, PyTorch-Direct also saves system power by 12.4% to 17.5% during training. To minimize programmer effort, we introduce a new "unified tensor" type along with necessary changes to the PyTorch memory allocator, dispatch logic, and placement rules. As a result, users need to change at most two lines of their PyTorch GNN training code for each tensor object to take advantage of PyTorch-Direct.