arXiv:2604.02651cs.LGcs.AI2026-04

提出无需通信的采样与四维并行,实现超大规模图神经网络训练。

Communication-free Sampling and 4D Hybrid Parallelism for Scalable Mini-batch GNN Training

  • 各设备独立采样构建子图,彻底消除采样通信开销。
  • 支持2048张GPU训练,在ogbn-products上比现有方法快3.5倍。
  • 适合超大规模图数据训练,尤其适合多卡集群部署。

图神经网络广泛应用于各类真实场景的图数据学习。处理极大规模图需分布式训练,而基于采样的小批量训练是主流并行方式。现有方法因采样成本高且数据并行扩展性差,存在显著性能瓶颈。本文提出ScaleGNN,一种支持可扩展小批量GNN训练的4D并行框架,融合无通信分布式采样、三维矩阵乘法(3D PMM)与数据并行。ScaleGNN引入统一顶点采样算法,使每个进程(GPU)可本地构建其子图批次,无需跨进程通信。3D PMM显著降低通信开销,使小批量训练可扩展至远超传统数据并行的GPU数量。此外,通过采样与训练重叠、低精度传输、核融合及通信-计算重叠等优化进一步减少开销。在五个图数据集上评估,表明ScaleGNN可在Perlmutter上实现2048张GPU的强扩展性,于Frontier达2048 GCDs,Tuolumne达1024张GPU。在Perlmutter上,相比最先进基线,ogbn-products任务端到端训练速度提升3.5倍。

原文摘要 · Abstract (English)

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.

图神经网络分布式训练并行计算大模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。