图神经网络在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小时即可达到最先进的性能。