CaPGNN：通过联合缓存和资源感知的图划分来优化并行图神经网络的训练

《Neurocomputing》：CaPGNN: Optimizing parallel graph neural network training with joint caching and resource-aware graph partitioning

提出CaPGNN框架，结合自适应缓存和资源感知图划分，减少多GPU全批量GNN训练的通信开销，实验显示效率提升18.98倍，通信成本降低99%。

　　

Xianfeng Song|Yi Zou|Zheng Shi

华南理工大学微电子学院，中国广东

摘要

图结构数据在现实世界中无处不在，图神经网络（GNN）由于能够直接处理这种不规则数据而在各个领域变得越来越流行。然而，随着数据规模的增加，GNN的效率会降低。尽管并行训练可以提高性能，但增加的通信成本往往抵消了这些优势。为了解决这个问题，本文介绍了一种名为CaPGNN的新型并行全批量GNN训练框架，该框架在单服务器上使用多GPU进行训练。首先，考虑到一个分区中的远程顶点数量通常大于或等于本地顶点数量，并且可能存在许多重复顶点，我们提出了一种联合自适应缓存算法，该算法利用CPU和GPU内存，结合轻量级缓存更新和预取技术来有效减少冗余通信成本。此外，考虑到GPU之间的计算和通信能力差异，我们提出了一种受图稀疏化启发的、考虑通信和计算的启发式图划分算法。我们还实现了一个管道来重叠计算和通信。广泛的实验表明，CaPGNN将训练效率提高了多达18.98倍，并将通信成本降低了多达99%，在某些情况下甚至提高了准确性。最后，我们将CaPGNN扩展到了多机器多GPU环境。代码可在https://github.com/songxf1024/CaPGNN获取。

引言

与图像和文本等结构化数据不同，图表示由顶点和边组成的多功能数据结构，能够有效捕捉不规则数据中的复杂关系。图神经网络（GNN）是专门为图数据设计的深度学习模型，能够直接处理图结构信息。因此，GNN在社交网络[33]、推荐系统[47]、知识图谱[52]、图像处理[30]和电子设计自动化[27]等领域得到了广泛应用。此外，GNN还可以与其他技术（如强化学习[21]和Transformers[50]）结合，以进一步扩大其应用范围和发展潜力。

随着数据规模的快速增长，传统的GNN训练方法面临可扩展性挑战。尽管一些方法试图通过操作符级别的优化[13]、[14]、[15]、[16]、[17]、[18]、[19]、[20]、[21]、[22]、[23]、[24]、[25]、[26]、[27]、[28]、[29]、[30]、[31]、[32]、[33]、[34]、[35]来提高计算效率，但大规模图的计算和内存需求往往超出单个设备的容量[51]。分布式GNN通过利用多个设备，将大型图划分为较小的子图进行并行处理，从而克服了这些限制，提高了训练效率。根据不同的计算资源，分布式GNN训练系统可以分为四种类型[32]：(a) 单机多GPU系统、(b) GPU集群系统、(c) CPU集群系统和(d) 无服务器和边缘设备系统。其中，单机多GPU系统由于其高性能、简单性和成本效益而在实践中被广泛使用。

分布式GNN训练方法主要包括全批量[7]和小批量[9]。小批量训练使用图采样来减少内存消耗并扩展到大型图。然而，它引入了计算开销，存在丢失全局信息的风险，并且缺乏理论上的收敛保证[14]。相比之下，全批量训练保留了全局图信息，提供了稳定的梯度，并实现了更高的准确性[55]，尽管代价是增加了通信和同步开销。

然而，图结构中的固有依赖性为高效的分布式GNN训练带来了额外的挑战[32]。例如，在以顶点为中心的分区下，一个晕圈顶点可能会出现在多个分区中，这需要跨分区通信来交换其特征，从而导致由于重复访问而产生的冗余通信。此外，大量的边界顶点和高维的顶点特征加剧了通信开销。此外，由于顶点分布不均，负载不平衡问题也会出现，现有的图划分方法未能解决这个问题，因为它们只关注平衡顶点数量而不考虑GPU性能差异。

为了解决这些挑战，我们提出了CaPGNN，这是一种高效的缓存和划分框架，用于在单服务器上使用多GPU进行并行全批量GNN训练。我们引入了一种新的联合缓存策略来减少与边界顶点相关的通信冗余，并提出了一种考虑GPU性能差异的资源感知图划分算法，旨在减少通信和同步开销。我们的主要贡献总结如下：•

为了全面了解使用CPU和GPU进行的大规模GNN训练问题，我们进行了系统的研究。我们分享了我们的分析发现和思考，特别是关于边界顶点（即晕圈顶点）对并行GNN训练的影响。此外，我们以三个主要研究观察结果的形式总结了我们的发现，这些结果激发了这项工作。

•

为了减少来自边界顶点的通信开销，我们提出了一种联合自适应缓存算法（JACA）。具体来说，我们利用CPU和GPU内存作为两级缓存，以最小化冗余的顶点通信，并优先缓存由分析模型评估的高重要性顶点。通过固定内存调度、轻量级缓存更新、预取和容错管道进一步优化了性能。

•

为了最小化通信开销和同步成本，我们提出了一种资源感知划分算法（RAPA）。它结合了直观的指标来模拟GPU的通信和计算能力，并考虑了内存可用性。受图稀疏化的启发，它在图划分过程中调整晕圈顶点的规模，以减少通信开销同时保持准确性。

•

为了全面评估所提出的CaPGNN，我们进行了广泛的实验。理论分析和实验结果都证实了CaPGNN的有效性。结果表明，CaPGNN显著提高了训练效率。我们进一步将CaPGNN扩展到了多机器多GPU环境，证明了其在分布式环境中的适用性。

本文的结构如下。第2节分析了相关工作。第3节解释了我们的动机。第4节详细描述了所提出的联合缓存策略和图划分算法。第5节描述了实验设置和比较研究。第6节我们分享了对于当前工作的局限性和未来探索方向的思考。最后，第7节总结了本文。

章节摘录

并行和分布式GNN训练

随着数据规模的指数级增长，传统GNN在处理能力上遇到了效率瓶颈。分布式GNN通过高效利用多机器的计算和存储资源，逐渐成为研究热点。ROC[16]通过图划分和内存管理的优化提高了系统性能和可扩展性，尽管这也增加了系统复杂性。SANCUS[28]在全局图训练中引入了历史嵌入以减少

背景和动机

在本节中，我们介绍了与并行全批量GNN训练相关的背景，并提出了一系列激发CaPGNN的关键观察结果。论文中的关键符号和符号在表2中进行了总结。

系统设计

我们提出了CaPGNN，这是一种用于并行全批量GNN训练的新型缓存和划分框架。CaPGNN减少了冗余通信，并尽可能确保了负载平衡。训练在配备多个GPU的单机上进行（即工作者），每个GPU负责训练一个子图。图7显示了整体架构。请注意，CaPGNN的设计原则可以很容易地扩展到分布式系统，我们在实验中展示了这一点。

实验

为了系统地评估所提出的CaPGNN的有效性，我们在几个广泛使用的公共数据集上进行了广泛的实验，并将其与现有的代表性方法进行了比较。我们的实验设计包括以下部分：首先介绍实验设置和数据集；5.2 顶点重叠比例对缓存命中率的影响，5.3 缓存容量和分区对缓存命中率的影响，5.4 缓存容量和分区对训练周期时间的影响，5.5

讨论

所提出的CaPGNN表现出色。然而，我们发现了一些需要改进的地方，下面将进行进一步讨论。

划分。目前，CaPGNN在划分过程中仅调整晕圈顶点，内部顶点不受影响。因此，当GPU之间的性能差异显著时，即使移除所有晕圈顶点也可能无法实现负载平衡。在未来的工作中，我们计划通过调整晕圈顶点来优化划分过程

结论

在本文中，我们提出了CaPGNN，这是一个优化框架，它结合了联合自适应缓存和资源感知图划分来提高并行全批量GNN训练的效率。通过利用其两个关键创新，即JACA（用于CPU和GPU之间的高效联合自适应缓存）和RAPA（用于有效的资源感知划分），CaPGNN显著减少了通信开销。理论分析和实验结果都证实了CaPGNN的有效性。

CRediT作者贡献声明

Xianfeng Song：撰写 – 审阅与编辑，撰写 – 原始草稿，可视化，验证，软件，资源，方法论，调查，形式分析，数据管理，概念化。Yi Zou：撰写 – 审阅与编辑，可视化，项目管理，资金获取。Zheng Shi：撰写 – 审阅与编辑。

利益冲突声明

作者声明他们没有已知的可能会影响本文报告工作的竞争性财务利益或个人关系。

致谢

这项研究部分得到了国家自然科学基金（项目编号U24B20151）和华南理工大学研究启动基金（项目编号K3200890）的支持。本文中表达的任何观点、发现、结论或建议均为作者的观点，并不一定反映资助机构的观点。

Xianfeng Song出生于中国浙江湖州。他于2015年获得江西南昌大学的学士学位。他目前正在华南理工大学微电子学院攻读博士学位。他的研究兴趣包括分布式计算、图神经网络和分布式存储。