无通信抽样与4维混合并行：可扩展小批量图神经网络训练
机器学习
2026-04-06 v1 人工智能
分布式、并行与集群计算
摘要
图神经网络（GNN）广泛应用于从各种实际场景中派生的图数据集上进行学习。极大图的学习需要分布式训练，小批量抽样是并行化GNN训练的流行方法。现有分布式小批量方法因抽样方法昂贵且数据并行扩展受限而存在显著性能瓶颈。本文提出ScaleGNN，这一用于可扩展小批量GNN训练的4维并行框架，结合无通信分布式抽样、3维矩阵乘法（PMM）和数据并行。ScaleGNN引入统一的顶点抽样算法，使每个进程（GPU设备）能够在无需进程间通信的情况下构建其本地小批量，即子图分区。3维PMM使小批量训练能够扩展至远超常规数据并行的GPU数量，通信开销显著降低。我们还提出了额外的优化措施，通过重叠抽样与训练、以更低精度发送数据、内核融合以及通信计算重叠来减少通信开销。我们在五个图数据集上对ScaleGNN进行评估，在Perlmutter上实现了2048个GPU的强扩展，在Frontier上实现了2048个GCD，在Tuolumne上实现了1024个GPU。 在Perlmutter上，ScaleGNN相对于SOTA基线在ogbn-products上实现了3.5倍端到端训练加速。
引用
@article{arxiv.2604.02651,
title = {Communication-free Sampling and 4D Hybrid Parallelism for Scalable Mini-batch GNN Training},
author = {Cunyang Wei and Siddharth Singh and Aishwarya Sarkar and Daniel Nichols and Tisha Patel and Aditya K. Ranjan and Sayan Ghosh and Ali Jannesari and Nathan R. Tallent and Abhinav Bhatele},
journal= {arXiv preprint arXiv:2604.02651},
year = {2026}
}