PaGraph由来自中科大李诚教授团队的林郅琦主导实现，该工作共发表在SoCC’20和TPDS’21上，代码已开源。

PaGraph的主要工作是在GNN系统中引入了特征缓存来加速训练并提供了一种简单可行的缓存策略；同时设计了一种图分区算法以支持对GPU的可扩展性。

背景知识

GNN

图是一种表示能力非常强大的数据结构，传统的文本、图像、视频的表示均可视作图的特例。图神经网络是传统神经网络在图拓扑数据上的扩展，在节点分类、链路预测、图生成领域拥有优秀的表现。

现实世界中的图往往规模庞大，例如ogbn-papers100M拥有111,059,956个节点和1,615,685,872条边，存储稀疏邻接矩阵所需的空间超过50GB，特征向量也占据了很大的体积。

minibatch

通常神经网络训练时是将所有训练数据送往GPU进行计算，全部的训练数据称之为batch，batch包含的样本数量为batch size，当batch size增大时，训练速度会变慢同时训练所占用的GPU内存也会增加（包括样本本身的内存占用和中间结果的内存占用），minibatch是一种扩展数据规模的有效训练方式。

方法

缓存

观察到在训练基于采样的GNN模型时，由于单个minibatch的数据量较小，往往不能完全利用GPU内存空间，因此我们可以将一部分特征提前保存在GPU内存中作为特征缓存。

缓存策略

由于图中节点与边连接的灵活性和复杂性，导致图的访存行为非常不规则，典型的动态缓存策略如FIFO、LRU在图数据集上的表现并不好；当缓存位于GPU内存时，动态缓存策略也有较高的管理开销。

现实世界中的图往往符合幂律分布，即少数节点占据了图中大量的边，缓存这些高度节点就可以获得较大的收益。PaGraph在预处理阶段将节点按出度大小降序排序，缓存排名靠前的节点。

但是，当图的倾斜度较小时，静态缓存策略的命中率也会降低。

缓存空间

在不影响训练的情况下，要尽可能地最大化缓存空间。使用minibatch训练时，由于minibatch的大小是固定的，训练产生的中间结果占用的显存也是接近的，因此只需要统计一次训练过程中的显存占用峰值，就可以得到可用于缓存的显存空间大小。

讨论

研究表明，当图的规模或训练时的batch size持续增大时，PaGraph的缓存性能会迅速下降，这是因为缓存的特征只占据了全图的一小部分，而每次迭代的batch size却增大了，导致命中率下降。

按照出度排序的静态缓存策略在图符合幂律分布以及使用典型的采样算法的情况下能保持较好的性能，但是对于一些倾斜度较小的图或者特殊的采样算法时性能会降低。

分区

动机

通常在多GPU场景的GNN训练中，在内存中只维护一份图拓扑数据，如果将上述缓存策略直接应用于该场景，命中率会随着GPU的增加而不断下降，这是因为所有GPU中都缓存了相同的特征数据。

为了解决该问题，PaGraph提出了一种数据并行训练方式，即将图划分为多个子图，每个GPU只训练自己的那一部分，这样做GPU就只需要缓存一个子图的特征，总体上来看扩展了缓存的特征的数量。

适用于这种数据并行的图划分算法需要满足两个要求：

• 负载均衡：每个分区的训练节点数量应该接近
• 数据隔离：尽可能避免不同GPU的跨分区数据访问

算法

假设需要将图划分为K个分区时，遍历训练集，每次迭代期间为每个训练节点分配一个K维向量，其中第i个元素代表将该训练节点分配给第i个分区的可能性，分数由如下公式计算：

 $$ socre_{\upsilon_t}^{(i)}=\mid TV_i\cap IN(V_t)\mid \cdot \frac{TV_{avg}-\mid TV_i\mid}{\mid PV_i\mid}mc^2 $$

其中，$TV_i$表示已经分配给分区i的训练节点集，$IN(V_t)$表示该训练节点的L跳邻居节点集，$TV_(avg)$表示最终每个分区的平均训练节点数，$PV_i$表示当前该分区已有的所有节点

算法流程如下：

• 初始化K个分区为空
• 遍历训练集的节点
• 对于每个节点获取其L跳邻居节点集
• 计算score
• 得到score最大的分区索引
• 将该训练节点和其L跳邻居节点集加入对应分区

优化

对于包含跨分区的边的查询，仍然需要将它们转发到图存储服务器，为了解决这个问题，PaGraph为每个分区引入了最少的冗余节点和边。训练节点也可能作为冗余节点被引入，但是这些节点并不会被训练。

为了进一步减少图存储的负担，PaGraph删除了那些在训练中没有贡献的节点和边，也就是说被删除的这些节点不会被任何训练节点在L跳内采样到。

讨论

PaGraph的设计针对单机多卡的分区算法的时间复杂度较大，对于大规模图并且分区数量较多的场景其时间开销难以接受。

不同的分区之间存在较多的冗余节点，这些冗余节点一部分来自于不同分区的训练节点的公共L跳邻居，另一部分来自为了减少边的跨分区访问而引入的冗余。

PaGraph在单机多卡场景下，每个GPU上训练完全独立的图分区，这忽略了GPU间通信技术，造成了闲置的总线带宽。

杂项

资源隔离

python由于GIL的存在，使用多线程来实现多卡并行计算性能很差，因此应该使用DDP方法，也就是多进程。

此外采样过程和数据加载都会争用CPU资源，采样和数据加载也应该使用单独的进程，并调整OpenMP配置来平衡它们之间的CPU资源

局部shuffing

全局洗牌一般收敛得更快，但可能成本较高

研究表明局部洗牌可以在收敛速度稍慢的情况下拥有良好的表现

分布式

其分区和缓存思想可以应用于分布式GNN训练