图神经网络CPU-GPU集群训练优化方案
在测试中,新方法的训练速度比前代技术快15到18倍。
图神经网络的应用挑战
图是表示数据的有用方式,因为它们捕获数据项之间的连接,而图神经网络(GNNs)是处理图的日益流行的方法。GNN在行业中的常见应用包括推荐、搜索和欺诈检测。
工业应用中使用的图通常非常庞大,具有数十亿个节点和数千亿甚至数万亿条边。在这种规模的图上训练GNN需要巨大的内存存储和计算能力,相应地需要较长的训练时间和较大的能源消耗。
DistDGLv2分布式训练方法
在新论文中,我们提出了一种分布式训练GNN的新方法,该方法同时使用CPU和GPU,优化了不同处理器类型的任务分配以最小化训练时间。
DistDGLv2具有三个主要组件:
- 分布式键值数据库(KVStore),用于存储节点/边特征和可学习嵌入;
- 分布式图存储,用于保存分区图以进行小批量采样;
- 一组训练器,在小批量上运行前向和后向计算以估计模型参数的梯度。
为了优化计算资源的使用并扩展到非常大的图,我们将这些组件在CPU和GPU之间划分。分布式KVStore和图存储使用CPU内存,CPU生成小批量。训练器将小批量数据读入GPU进行小批量计算。
关键技术优化
加速DistDGLv2中小批量训练的关键是有效地将小批量从CPU移动到GPU。为此,DistDGLv2部署了三种策略:
首先,它使用METIS图分区程序生成具有最小边切割的图分区,并将数据与计算共置以减少网络通信;
它构建异步小批量训练流水线,在所有硬件中重叠计算和数据移动;
它将尽可能多的计算移动到GPU,以利用GPU的计算能力。
流水线优化
DistDGLv2通过将小批量流水线分成七个阶段来实现第二和第三种策略,其中五个阶段帮助准备小批量。我们尽可能多地将阶段保持在GPU上以利用GPU的计算能力,同时将小批量采样阶段放在另一个线程的CPU中。这使我们能够重叠GPU中的小批量计算和CPU中的小批量采样。
除此之外,我们进一步重叠网络通信和CPU计算。我们让采样流水线"前瞻"并同时采样多个小批量。因此,当生成一个小批量时,当给定的CPU等待远程邻居采样(从另一台机器)或特征复制(到GPU)时,它可以移动到另一个小批量以在本地采样邻居或复制数据。通过这种方式,我们可以有效地隐藏网络通信延迟。
性能比较
通过这些优化,DistDGLv2可以在CPU和GPU集群中有效地执行分布式GNN训练。我们在具有各种GNN工作负载的g4dn.metal实例集群上展示了DistDGLv2的效率。DistDGLv2相对于仅使用CPU的方法的性能表明,对于大规模图上的分布式GNN小批量训练,GPU可能比CPU更有效。
研究人员还提出了对GNN模型使用全图训练的方法。该方法在整个图上运行前向和后向计算。我们在相同的图数据集和相同的硬件上比较了小批量训练和全图训练。我们表明,小批量训练训练GNN模型要高效得多,并且图越大,速度差距越大。
在基于OGBN-papers100M数据集构建的图上(具有1亿个节点),小批量训练比全图训练快约100倍。经过六天的训练,全图训练仍然无法达到与小批量训练相同的精度,而小批量训练在相同的CPU上仅需1.5小时即可达到最先进的性能。