摘要
arXiv:2505.04083v1 类型: cross
摘要: 图神经网络已成为一类能够利用现实世界图的连接性和结构来学习节点之间复杂属性和关系的神经网络。许多现实世界中的图由于其巨大的规模而超出了GPU的内存容量,使用图神经网络对其进行处理需要采用像小批量采样这样的技术来进行扩展。然而,在某些情况下,这会导致准确性的下降,而采样和数据从CPU传输到GPU也可能减慢训练速度。另一方面,分布式全图训练由于图的不规则结构而导致高通信开销和负载不均衡。我们提出了一种名为Plexus的三维(3D)并行方法,该方法解决了这些问题,并可扩展到包含数十亿条边的图。此外,我们引入了平衡负载的排列方案,并开发了预测最佳3D配置的性能模型。我们在多个图数据集上评估了Plexus,并在Perlmutter机器上的2048个GPU(机器的33%)和Frontier机器上的2048个GCD上展示了扩展结果。Plexus在Perlmutter上实现了超过2.3倍至12.5倍的加速,而在Frontier上实现了5.2倍至54.2倍的加速,显著减少了求解时间。