Plexus：通过三维并行全图 GNN 训练驯服十亿边图
机器学习
2025-10-30 v2 人工智能
分布式、并行与集群计算
摘要
图神经网络（GNN）利用真实世界图的连接性和结构来学习节点间复杂的属性和关系。许多真实世界图因其庞大的规模而超出了单个 GPU 的存储容量，在此类图上训练 GNN 需要使用小批量采样等技术进行扩展。分布式全图训练的替代方法则因图的不规则结构而面临高通信开销和负载不均衡的问题。我们提出了一种用于全图训练的三维（3D）并行方法，以解决这些问题并扩展到十亿边规模的图。此外，我们还引入了诸如用于负载均衡的双重排列方案，以及用于预测我们并行实现——Plexus——最优 3D 配置的性能模型等优化。我们在六个不同的图数据集上评估了 Plexus，并展示了在 Perlmutter 上多达 2048 个 GPU 和 Frontier 上多达 1024 个 GPU 的扩展结果。Plexus 实现了前所未有的加速，比先前的最先进技术快 2.3-12.5 倍，并将求解时间在 Perlmutter 上减少了 5.2-8.7 倍，在 Frontier 上减少了 7.0-54.2 倍。
引用
@article{arxiv.2505.04083,
title = {Plexus: Taming Billion-edge Graphs with 3D Parallel Full-graph GNN Training},
author = {Aditya K. Ranjan and Siddharth Singh and Cunyang Wei and Abhinav Bhatele},
journal= {arXiv preprint arXiv:2505.04083},
year = {2025}
}