本文介绍一种可以高效地处理大规模图神经网络分布式训练框架。该论文成果论文 NeutronStar: Distributed GNN Training with Hybrid Dependency Management已经被SIGMOD2022接受为 Data Management Track。

背景

图神经网络算法结合了深度神经网络技术与图广播操作，因其强大的表达能力而广泛应用于自然语言处理，知识图谱，和生物信息学等领域。考虑到真实世界图数据的大规模，我们需要构建基于大规模并行的分布式系统来处理GNN计算。

 

在DNN训练中，输入的训练数据之间没有依赖关系。例如图片分类任务中， 输入图片数据集可以被均匀地划分到不同的子集中做并行训练，不同训练子集之间不需要交换信息（数据并行策略）。与之相对应，图神经网络计算需要处理复杂的顶点依赖关系（即一个顶点和其所有直接入邻居之间的依赖关系），由于神经网络往往由多个图广播和NN计算层串联，因此图神经网络系统需要在训练过程中处理复杂的多层依赖关系。

 

目前涌现出许多分布式图神经网络系统，如AliGraph，DistDGL，ROC，DistGNN等等。这些系统根据其依赖管理方式可以分成两类：基于依赖缓存的系统如 AliGraph，DistDGL等等。和基于依赖通信的系统如DistGNN，ROC等等。基于依赖缓存的方法将训练点均匀划分到集群中，然后在将顶点所依赖的训练子图和顶点标签存储到本地。在训练过程中，顶点不在需要动态地从其他机器获取到所依赖的顶点属性和顶点表示。而基于依赖通信的方法将训练顶点和其直接依赖均匀地划分。在训练过程中，不同子图通过通信动态地获取顶点依赖的feature和representation。现有系统有效地支持了大规模图神经网络的分布式训练，然而这两种方法在性能和实用性方面仍然受到诸多问题：

 

● 基于依赖缓存的系统，前向计算与反向计算都可以在单个机器内执行，因此可以较好地匹配基于数据并行的深度学习系统。然而，由于不同顶点可以存在相同的数据依赖，不同机器上的训练子图可能存在重叠，因此造成冗余计算问题。

● 基于依赖通信的系统，通过通信处理机器之间的依赖关系，这样一来，冗余计算问题被去除，然而其通信代价高昂，且由于计算图被切断，我们无法直接使用现有自动梯度计算库实现自动梯度反向传播操作。

混合依赖图处理策略

在一系列实验之后，我们发现这两种方法在性能上各有优劣，其性能受数据集，模型设置，和环境配置多方面影响。

 

因此，我们提出一种混合依赖管理策略Hybrid，与传统使用单一策略的方法不同，Hybrid合并了基于依赖缓存和基于依赖通信的策略来最大化性能。我们的方法的关键在于基于给定输入图，执行环境和模型设置，我们评估每一个依赖点在基于依赖缓存和基于依赖通信策略下的开销，然后选择最合适的处理策略。

 

NeutronStar

基于此，我们提出了NeuTronStar，一个基于GPU加速的分布式图神经网络系统，在混合依赖处理策略基础之上。NeutronStar解耦了NN计算与图依赖处理操作，并提出一种灵活的自动微分框架。这样一来。NeutronStar可以使用现有GPU-优化的自动微分库来实现NN计算。此外，NeutronStar还实现了一系列优化来高效地管理依赖并调度计算任务。如chunk-based图划分，环形任务调度，无锁的并行消息队列，和通信计算重叠等等。 此外NeutronStar还提供pythonAPI以提供较好的易用性。

 

实验评估

我们在一个16节点的阿里云GPU集群上评估了NeutronStar的性能。实验系统包括DistDGL，ROC，和基于NeutronStar代码库实现的基于依赖缓存（DepCache）的和基于依赖通信(DepComm)的分布式训练系统。

 

上图显示了在Google,Pokec, liveJournal, reddit, Orkut, Wiki, 和Twitter 数据集上执行GCN，GIN，和GAT算法的性能。实验结果表明，与其他系统和策略相比NeutronStar展示出了非常优异的性能。具体地，是DistDGL的1.83到14.25倍，是ROC的1.81到5.29倍，是DepCache的2.03到15.02倍，是DepComm的1.51到2.21倍。

 

总结

我们设计并实现了NeutronStar，一个支持依赖缓存/依赖通信混合依赖管理，灵活自动微分框架，丰富GPU-CPU异构优化和通信优化的分布式图神经网络训练系统。与现有工作相比，NeutronStar具有较高的性能较好的可扩展性和较为均衡的通信计算资源利用。