基于高效采样算法的时序图神经网络系统

现有的图数据规模极大，导致时序图神经网络的训练需要格外长的时间，因此使用多GPU进行训练变得成为尤为重要，如何有效地将多GPU用于时序图神经网络训练成为一个非常重要的研究议题。本文提供了两种方式来提高多GPU训练的性能，包括局部性感知的数据划分策略以及高效的任务调度策略。

首先，本文使用高效的数据划分策略将图数据以及点的特征向量进行有效的划分，来降低GPU之间额外的通信开销。点的特征向量占据很大的存储空间，因此我们将点的特征向量按照三个维度进行划分并且存储在不同的GPU中。

如图1所示，首先我们按照时序图的时间维度进行切分。在不同的时间范围中，将该时间范围内的点数据按照源点和汇点进行二维划分，其中每个GPU存储一段区间内源点的特征向量。对于边数据的存储，我们将边数据直接存储在CPU memory中。在进行训练时再将对应的边数据加载进相应的GPU中。

图1. 局部性感知的数据划分策略
由于不同的数据存储在不同GPU中，如何进行高效的任务调度来降低通信开销将成为一个至关重要的问题。为解决该问题，我们提出了一种高效的任务调度策略：

首先我们的任务调度将存储在不同GPU上的数据进行调度，从而充分利用GPU之间NVLink高速通信的优势。

在进行训练时，由于每个GPU存储一段区间内的点特征向量，每个GPU需要传输当前边数据对应的汇点的特征向量。在8块V100GPU中，部分GPU之间存在两条NVLink高速链路，我们将充分利用这些NVLink链路来传输点的特征向量，从而提升传输效率。

以8块GPU为例，我们一共进行4轮的点数据交互。