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.

翻译：图神经网络（GNN）被广泛用于从各种现实场景中提取的图数据集上进行学习。针对超大规模图的学习需要分布式训练，而结合采样的小批量方法常被用于并行化GNN训练。现有分布式小批量方法因昂贵的采样技术以及数据并行下的有限扩展性，存在显著性能瓶颈。本文提出ScaleGNN——一种面向可扩展小批量GNN训练的4D并行框架，该框架融合了无通信分布式采样、3D并行矩阵乘法（PMM）与数据并行。ScaleGNN引入均匀顶点采样算法，使得每个进程（GPU设备）能够在无需任何进程间通信的情况下构建本地小批量（即子图分区）。3D PMM可使小批量训练扩展到远高于普通数据并行所能支持的GPU数量，同时显著降低通信开销。我们还提出了额外优化措施，包括将采样与训练重叠、通过低精度数据发送降低通信开销、内核融合以及通信-计算重叠。我们在五个图数据集上评估了ScaleGNN，并在Perlmutter上实现至多2048个GPU、Frontier上2048个GCD以及Tuolumne上1024个GPU的强扩展。在Perlmutter上，ScaleGNN相较于ogbn-products数据集上的SOTA基线实现了3.5倍的端到端训练加速。