扩展图神经网络训练的CPU-GPU集群方法

在测试中，新方法的速度达到前代方法的15到18倍。

图是一种有用的数据表示形式，能够捕捉数据项之间的连接关系，而图神经网络（GNNs）正成为处理图的日益流行的方法。GNN的常见行业应用包括推荐系统、搜索和欺诈检测。工业应用中使用的图通常规模巨大，包含数十亿个节点和数千亿甚至数万亿条边。在此规模图上训练GNN需要巨大的内存存储和计算能力，相应地导致训练时间长和能源消耗大。

在一篇论文中，描述了一种新的分布式训练GNN的方法，该方法同时使用CPU和GPU，优化不同处理器类型的任务分配以最小化训练时间。在测试中，该方法——DistDGLv2——在相同硬件上比另一种分布式GNN训练框架Euler快18倍。DistDGLv2在相同规模的集群中，也比分布式CPU训练快15倍。

图神经网络

在GNN设置中，图节点通常表示对象，图边表示对象之间的关系。节点和边都可能具有相关特征——例如对象属性或对象间关系类型的数据。

对于图中的每个节点，GNN生成一个向量表示（嵌入），编码有关节点及其邻域的信息——通常是一跳或两跳邻域，但有时是更大的区域。在工业应用中常见的大规模图中，考虑节点的所有一跳邻居可能非常耗时，更不用说更远的邻居。因此，在生成节点嵌入时，GNN通常使用从目标节点邻域采样的迷你批次节点。

有许多关于迷你批次采样的研究工作——例如，在KDD 2021上提出的全局邻居采样技术。在新论文中，实现了一种由GraphSage提出的流行迷你批次采样算法，如下图所示。它首先采样目标节点（如蓝色节点），然后采样它们的邻居节点（如红色节点和橙色节点）。然而，DistDGLv2具有实现其他采样算法的灵活性。

DistDGLv2

DistDGLv2具有三个主要组件：

• 分布式键值数据库（KVStore），用于存储节点/边特征和可学习嵌入；
• 分布式图存储，用于保存分区图以进行迷你批次采样；
• 一组训练器，在迷你批次上运行前向和后向计算以估计模型参数的梯度。

为了优化计算资源的使用并扩展到非常大的图，将这些组件在CPU和GPU之间分配。分布式KVStore和图存储使用CPU内存，CPU生成迷你批次。训练器将迷你批次数据读入GPU进行迷你批次计算。

在DistDGLv2中加速迷你批次训练的关键是高效地将迷你批次从CPU移动到GPU。为此，DistDGLv2部署了三种策略：

首先，使用METIS图分区程序（由某机构高级首席科学家共同开发）生成具有最小边切割的图分区，并将数据与计算共置以减少网络通信； 构建异步迷你批次训练管道，以重叠所有硬件中的计算和数据移动； 尽可能多地将计算移动到GPU，以利用GPU的计算能力。

为了将数据与计算共置，DistDGLv2在同一组机器上运行KVStore服务器、分布式图存储服务器和训练器。当加载图分区时，其节点和边特征进入KVStore，图结构进入图存储服务器。每个训练器被分配一个训练集，其中大多数训练节点和边属于分配给同一台机器的图分区。通过这种方式，在迷你批次训练期间，与迷你批次相关的大多数数据将来自本地机器。

DistDGLv2通过将迷你批次管道拆分为七个阶段来实现第二和第三种策略，其中五个阶段帮助准备迷你批次。尽可能多地将阶段保持在GPU上以利用GPU的计算能力，同时将迷你批次采样阶段放在CPU的另一个线程中。这使得可以重叠GPU中的迷你批次计算和CPU中的迷你批次采样。

如下图所示，在GPU中运行最后四个阶段；其中一些阶段仍然参与迷你批次准备。除此之外，进一步重叠网络通信和CPU计算。让采样管道“前瞻”并同时采样多个迷你批次。因此，当生成一个迷你批次时，当给定的CPU等待远程邻居采样（从另一台机器）或特征复制（到GPU）时，它可以移动到另一个迷你批次以在本地采样邻居或复制数据。通过这种方式，可以有效地隐藏网络通信延迟。

通过这些优化，DistDGLv2可以在CPU和GPU集群中有效执行分布式GNN训练。在g4dn.metal实例集群上使用各种GNN工作负载展示了DistDGLv2的效率。DistDGLv2相对于纯CPU方法的性能表明，对于大规模图上的分布式GNN迷你批次训练，GPU比CPU更有效。

研究人员还提出了使用全图训练GNN模型的方法。该方法在整个图上运行前向和后向计算。在相同图数据集和相同硬件上比较了迷你批次训练和全图训练。表明迷你批次训练训练GNN模型效率更高，并且随着图规模的增大，速度差距变得更大。在基于OGBN-papers100M数据集构建的图上（具有1亿个节点），迷你批次训练速度约为全图训练的100倍。经过六天的训练，全图训练仍然无法达到与迷你批次训练相同的准确性，而迷你批次训练在相同CPU上仅需1.5小时即可达到最先进的性能。