Communication-free Sampling and 4D Hybrid Parallelism for Scalable Mini-batch GNN Training
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866911564426117120 |
|---|---|
| author | Wei, Cunyang Singh, Siddharth Sarkar, Aishwarya Nichols, Daniel Patel, Tisha Ranjan, Aditya K. Ghosh, Sayan Jannesari, Ali Tallent, Nathan R. Bhatele, Abhinav |
| author_facet | Wei, Cunyang Singh, Siddharth Sarkar, Aishwarya Nichols, Daniel Patel, Tisha Ranjan, Aditya K. Ghosh, Sayan Jannesari, Ali Tallent, Nathan R. Bhatele, Abhinav |
| contents | Graph neural networks (GNNs) are widely used for learning on graph datasets derived from various real-world scenarios. Learning from extremely large graphs requires distributed training, and mini-batching with sampling is a popular approach for parallelizing GNN training. Existing distributed mini-batch approaches have significant performance bottlenecks due to expensive sampling methods and limited scaling when using data parallelism. In this work, we present ScaleGNN, a 4D parallel framework for scalable mini-batch GNN training that combines communication-free distributed sampling, 3D parallel matrix multiplication (PMM), and data parallelism. ScaleGNN introduces a uniform vertex sampling algorithm, enabling each process (GPU device) to construct its local mini-batch, i.e., subgraph partitions without any inter-process communication. 3D PMM enables scaling mini-batch training to much larger GPU counts than vanilla data parallelism with significantly lower communication overheads. We also present additional optimizations to overlap sampling with training, reduce communication overhead by sending data in lower precision, kernel fusion, and communication-computation overlap. We evaluate ScaleGNN on five graph datasets and demonstrate strong scaling up to 2048 GPUs on Perlmutter, 2048 GCDs on Frontier, and 1024 GPUs on Tuolumne. On Perlmutter, ScaleGNN achieves 3.5x end-to-end training speedup over the SOTA baseline on ogbn-products. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2604_02651 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Communication-free Sampling and 4D Hybrid Parallelism for Scalable Mini-batch GNN Training Wei, Cunyang Singh, Siddharth Sarkar, Aishwarya Nichols, Daniel Patel, Tisha Ranjan, Aditya K. Ghosh, Sayan Jannesari, Ali Tallent, Nathan R. Bhatele, Abhinav Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing Graph neural networks (GNNs) are widely used for learning on graph datasets derived from various real-world scenarios. Learning from extremely large graphs requires distributed training, and mini-batching with sampling is a popular approach for parallelizing GNN training. Existing distributed mini-batch approaches have significant performance bottlenecks due to expensive sampling methods and limited scaling when using data parallelism. In this work, we present ScaleGNN, a 4D parallel framework for scalable mini-batch GNN training that combines communication-free distributed sampling, 3D parallel matrix multiplication (PMM), and data parallelism. ScaleGNN introduces a uniform vertex sampling algorithm, enabling each process (GPU device) to construct its local mini-batch, i.e., subgraph partitions without any inter-process communication. 3D PMM enables scaling mini-batch training to much larger GPU counts than vanilla data parallelism with significantly lower communication overheads. We also present additional optimizations to overlap sampling with training, reduce communication overhead by sending data in lower precision, kernel fusion, and communication-computation overlap. We evaluate ScaleGNN on five graph datasets and demonstrate strong scaling up to 2048 GPUs on Perlmutter, 2048 GCDs on Frontier, and 1024 GPUs on Tuolumne. On Perlmutter, ScaleGNN achieves 3.5x end-to-end training speedup over the SOTA baseline on ogbn-products. |
| title | Communication-free Sampling and 4D Hybrid Parallelism for Scalable Mini-batch GNN Training |
| topic | Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing |
| url | https://arxiv.org/abs/2604.02651 |