图神经网络(GNNs)是一种基于图的深度学习模型,已成功应用于许多领域。尽管GNN具有一定的有效性,但GNN要有效地扩展到大型图仍然是一个挑战。作为一种补救措施,分布式计算能够提供丰富的计算资源,成为训练大规模GNN的一种很有前途的解决方案。然而,图结构的依赖性增加了实现高效分布式GNN训练的难度,这将受到海量通信和工作量不平衡的影响。近年来,人们在分布式GNN训练方面做了很多工作,提出了一系列的训练算法和系统。然而,从图处理到分布式执行的优化技术还缺乏系统的综述。本文分析了分布式GNN训练中的三个主要挑战,即大量特征通信、模型精度损失和工作负载不平衡。然后,我们介绍了分布式GNN训练优化技术的一种新的分类法,以解决上述挑战。新的分类法将现有技术分为四类:GNN数据分区、GNN批处理生成、GNN执行模型和GNN通信协议。我们仔细讨论了每个类别中的技术。最后,我们分别总结了现有的多GPU、GPU集群和CPU集群分布式GNN系统,并对可扩展GNN的未来发展方向进行了讨论。
https://arxiv.org/abs/2211.00216
1.『概述』
图神经网络(GNN)是处理图建模问题的强大工具,已被广泛应用于各种应用中,包括社交网络(如社交垃圾邮件检测[87,112],社会网络分析[91]),生物信息学(如蛋白质界面预测[35],疾病-基因关联[83]),药物发现[11,66],流量预测[61],医疗保健[2,21],推荐[32,49,55,111],自然语言处理[134,145]和其他[26,85,133,141,144]。通过将图结构信息融入到深度学习模型中,GNNs可以取得明显优于传统机器学习和数据挖掘方法的效果。GNN模型通常包含多图卷积层,其中每个顶点聚合其邻居的最新状态,更新顶点的状态,并对更新的顶点状态应用神经网络(即转换)。以传统的图卷积网络(GCN)为例,在每一层中,顶点使用求和函数聚合相邻的状态和自己的状态,然后应用单层MLP转换新状态。如果层数为𝐿,这样的过程重复𝐿次。在𝐿th层中生成的顶点状态供下游任务使用,如节点分类、链接预测等。近年来,许多研究工作在图神经网络模型的设计方面取得了显著进展。著名的模型包括GCN[108]、GraphSAGE[46]、GAT[95]、GIN[118]和许多其他特定于应用程序的GNN模型[128,130]。到目前为止,已有数十篇综述综述了GNN模型[114,116,135,147]。另一方面,为了高效地开发不同的GNN模型,基于各种深度学习库提出了许多面向GNN的框架[8,14,34,44,68,104]。许多新的优化被提出来加速GNN训练,包括GNN计算内核[19,38,52,53,81,93,129],高效的编程模型[51,113,117],以及充分利用新硬件[18,42,124,148]。然而,这些框架和优化主要集中在单个机器上训练GNN,而没有过多关注输入图的可扩展性。
如今,大规模图神经网络[62,70]由于大规模大图数据的普遍存在而成为一个热门话题。拥有数十亿个顶点和数万亿条边的图形是很常见的,比如新浪微博、微信、Twitter和Facebook等社交网络。然而,现有的大多数GNN模型只在小图数据集上进行测试,无法处理大图数据集[50]或效率低下。这是因为之前的GNN模型比较复杂,在处理大型图时需要大量的计算资源。通过设计可扩展的GNN模型,实现了大规模的图神经网络。他们使用简化[36,47,110],量化[4,33,54,71,90,100,101,107,139],取样[20,123,125]和蒸馏[24,121,131]来设计高效的模型。另一种方法是将分布式计算应用于GNN训练。因为在处理大型图形时,单个设备(如GPU)有限的内存和计算资源成为大规模GNN训练的瓶颈,而分布式计算提供更多的计算资源(如多GPU、CPU集群等)来提高训练效率。在设计高效的分布式GNN训练框架(如PipeGCN [99], BNS-GCN[98])和分布式GNN训练系统(如AliGraph [150], DistDGL [142], ROC[58])方面已经做出了许多努力。在这个综述中,我们关注的工作与分布式计算的大图神经网络,又称分布式GNN训练。它将整个模型训练的工作量分配给一组工人,所有工人并行处理工作量。然而,由于GNN的数据依赖性,将现有的分布式机器学习方法[96,103]应用于GNN并非易事,许多新的技术被提出来优化分布式GNN的训练管道。尽管有很多关于GNN模型的综述[114,135,147],但据我们所知,很少有人系统地回顾分布式GNN训练的技术。最近,Besta等[9]只回顾了GNN的并行计算范式,Abadal[1]从算法到硬件加速器考察了GNN计算。
为了清晰地组织分布式GNN训练的技术,我们介绍了一种通用的分布式GNN训练管道,该管道由数据划分、GNN模型优化和梯度聚合三个阶段组成。前两个阶段涉及GNN特定的执行逻辑,包括图处理和图聚合。最后一个阶段与传统的分布式机器学习相同。在这个通用分布式GNN训练管道的背景下,我们讨论了分布式GNN训练的三个主要挑战,这些挑战是由图数据中的数据依赖性引起的,需要新的分布式GNN特定技术。为了帮助读者更好地理解解决上述挑战的各种优化技术,我们介绍了一种新的分类法,将这些技术分为四个正交的类别:GNN数据分区、GNN批处理生成、GNN执行模型和GNN通信协议。该分类法不仅涵盖了小批分布式GNN训练和全图分布式GNN训练中使用的优化技术,而且还讨论了从图处理到模型执行的技术。我们仔细回顾了每个类别中的现有技术,然后描述了23个来自业界或学术界的代表性分布式GNN系统和框架。最后,简要讨论了分布式GNN训练的未来发展方向。
这项综述的贡献如下:
- 这是第一个专注于高效分布式GNN训练优化技术的综述,它帮助研究人员快速了解分布式GNN训练的现状。
- 通过考虑端到端分布式GNN训练的生命周期,我们引入了分布式GNN训练技术的一种新的分类法。在高层次上,新的分类法由四个正交类别组成,分别是GNN数据分区、GNN批处理生成、GNN执行模型和GNN通信协议。
- 我们为新分类法中的每个类别提供了详细和全面的技术总结。
- 我们回顾了23个具有代表性的分布式GNN训练系统和框架,从工业界到学术界。
- 讨论分布式GNN训练的未来发展方向。
2.『分布式GNN训练和挑战』
为了更好地理解端到端分布式GNN训练的一般工作流程,我们将训练管道分为数据划分、GNN模型优化和梯度聚合三个阶段。图1可视化了端到端分布式GNN训练工作流的高级抽象。
- 数据分区。这是一个支持分布式训练的预处理阶段。它将输入数据(即图和特征)分布到一组工作者中。考虑到GNN训练数据的依赖性,数据划分阶段比传统分布式机器学习中的数据划分阶段更加复杂。如图1所示,分区数据(即子图)之间的跨工作者边暗示了数据依赖性。如果我们承认分区之间的数据依赖性,分布式训练效率会因通信而降低;如果我们简单地忽略数据依赖性,模型的准确性就会被破坏。因此,数据分区是端到端分布式GNN训练效率的关键阶段。
- GNN模型优化。这是分布式GNN训练的核心阶段,执行GNN模型的训练逻辑(即前向计算和后向计算)。我们进一步将这一阶段分为计算图生成、执行模型和通信协议。首先,每个worker根据分块的输入图和特征生成计算图,然后用计算图执行GNN模型,计算损失和梯度。但由于数据的依赖性,计算图的生成和GNN模型的执行与传统的深度学习模型有很大的不同。在不访问远程输入数据的情况下,小批量训练策略的计算图可能无法正确生成。执行模型涉及GNN模型的𝑘-layer图聚合,聚合呈现不规则的数据访问模式。在分布式全图训练中,每一层的图聚合需要通过通信协议访问顶点远程邻居的隐藏特征,还需要考虑层与层之间的同步模式。因此,与传统的分布式机器学习相比,由于GNN的数据依赖性,计算图的生成和模型的执行变得更加复杂。
- 梯度聚合。该阶段负责聚合最新的局部梯度,获得全局梯度,更新模型参数。在GNN环境下,模型规模通常较小,模型更新逻辑与其他机器学习模型相同。经典分布式机器学习中已有的梯度聚合技术可以直接应用于分布式GNN训练。
为了解决上述问题,实现分布式GNN训练并优化训练效率,近年来提出了许多新的技术。大多数现有的工作都是从系统或框架的角度来展示他们自己的技术贡献,而没有从技术的角度来很好地回顾相关的工作。本文介绍了一种新的分类法,以端到端分布式训练管道中的阶段为基础,组织分布式gnn特定技术。通过这样的设计,我们将分布式GNN训练管道中相同阶段的相似技术组织在一起,帮助读者充分理解分布式GNN训练中不同阶段的现有解决方案。根据以往的实证研究,由于数据依赖性,分布式GNN训练的瓶颈一般出现在数据分区和在建的GNN模型优化阶段,而针对GNN的新技术就在这两个阶段。此外,不同的训练策略(如小批量训练、全图训练)带来不同的工作负载模式,导致模型优化中使用不同的优化技术。例如,模型优化阶段的计算图生成对于小批训练很重要,而通信协议对于全图训练很重要。因此,我们的新分类法将特定于分布式GNN的技术分为四类(即GNN数据分区、GNN批处理生成、GNN执行模型和GNN通信协议),如图2所示。下面,我们将介绍每个类别的概述。
3.『GNN数据分区』
在本节中,我们回顾了分布式GNN训练中现有的GNN数据分区技术。图3描述了这些技术的概述。考虑到图和特征是GNN中两种典型的数据类型,我们将划分方法分为图划分和特征划分。优化目标是工作量平衡、通信和计算最小化,旨在解决挑战#1和#3。此外,成本模型是捕获GNN工作负载特征的另一个关键组件。在下面,我们首先提出了代价模型,然后分别讨论了图划分和特征划分。
4.『GNN批量生成』
小批GNN训练是将GNN扩展为大型图的常用方法。图采样实际上是在独立模式下生成小批量的工具。到目前为止,已经提出了许多基于采样的图神经网络[17,23,46,56,125,151],根据采样方法的不同类型,可以分为顶点采样、层采样和子图采样。不同的批生成方法会影响训练效率和训练精度。为了避免图采样成为瓶颈,人们对高效GNN数据加载器进行了一些探索[5,25,27,76]。在小批分布式GNN训练中,数据依赖性为批生成过程带来了大量的通信。为了提高分布式环境下的训练效率,针对分布式训练提出了几种新的GNN批生成技术,并解决了挑战#1和挑战#2。如图4所示,一种解决方案是通过分布式抽样生成小批处理,另一种是直接使用本地分区(或子图)作为小批处理。
5.『GNN执行模型』
图1中的GNN模型优化阶段包括计算图生成和计算图执行两个子阶段。执行模型负责对子阶段的操作进行调度,以达到较高的训练效率。这两个子阶段针对不同的训练范式有不同的运算符。在分布式小批训练中,采样和特征提取操作影响训练效率,计算图生成成本高。而对于分布式全图训练,由于工作者之间的数据依赖性,计算图的执行比较复杂。下面,我们分别讨论了计算图生成和计算图执行的执行模型。
6.『结论』
分布式GNN训练是将GNN模型扩展到大图的成功方法之一。在本综述中,我们系统地回顾了现有的分布式GNN训练技术,从图数据处理到分布式模型执行,涵盖了端到端分布式GNN训练的生命周期。我们将分布式GNN训练管道分为三个阶段,其中数据划分阶段和GNN模型优化阶段是影响GNN训练效率的两个关键阶段。为了清晰地组织优化这两个关键阶段的新技术贡献,我们提出了一个新的分类法,由四个正交的类别组成:GNN数据分区、GNN批处理生成、GNN执行模型和GNN通信协议。在GNN数据分区类别中,我们描述了用于分布式GNN训练的数据分区技术;在GNN批生成类中,我们提出了用于小批分布式GNN训练的快速GNN批生成技术;在GNN执行模型方面,讨论了计算图生成和计算图执行所使用的执行模型;在GNN通信协议类别中,我们讨论了分布式GNN训练的同步和异步协议。在仔细回顾了这四类技术之后,我们分别总结了现有代表性的多GPU分布式GNN系统、GPU集群分布式GNN系统和CPU集群分布式GNN系统,并对未来优化大规模GNN训练的方向进行了讨论。