另一篇:【阅读】A Comprehensive Survey on Distributed Training of Graph Neural Networks
文中,涉及到公式部分的翻译不准确,请看对应原文。
摘要
图神经网络(GNNs)是一种在图上学习的深度学习模型,并已成功应用于许多领域。尽管 GNN 有效,但 GNN 有效地扩展到大型图仍然具有挑战性。作为一种补救措施,分布式计算成为训练大规模 GNN 的一种有前途的解决方案,因为它能够提供丰富的计算资源。然而,图结构的依赖性增加了实现高效分布式 GNN 训练的难度,导致大量通信和工作负载不平衡。近年来,在分布式 GNN 训练方面做出了很多努力,并提出了一系列训练算法和系统。然而,缺乏对从图处理到分布式执行的优化技术的系统回顾。在本次调查中,我们分析了分布式 GNN 训练的三大挑战,即海量特征通信、模型精度损失和工作负载不平衡。然后,我们为分布式 GNN 训练中的优化技术引入了一种新的分类法,以应对上述挑战。新分类法将现有技术分为四类,即 GNN 数据分区、GNN 批处理生成、GNN 执行模型和 GNN 通信协议。我们仔细讨论了每个类别中的技术。最后,我们分别总结了用于多 GPU、GPU 集群和 CPU 集群的现有分布式 GNN 系统,并讨论了可扩展 GNN 的未来发展方向。
简介
GNN 是处理图建模问题的强大工具,并已广泛应用于各种应用,包括社交网络(例如,社交垃圾邮件发送者检测 [87、112]、社交网络分析 [91])、生物信息学(例如,蛋白质界面预测 [35]、疾病-基因关联 [83])、药物发现 [11、66]、流量预测[61]、医疗保健 [2、21]、推荐 [32、49、 55、111]、自然语言处理 [134、145] 和其他 [26、85、133、141、144]。通过将图结构信息集成到深度学习模型中,GNN 可以获得比传统机器学习和数据挖掘方法明显更好的结果。
GNN 模型通常包含多图卷积层,其中每个顶点聚合其邻居的最新状态,更新顶点的状态,并将神经网络(即转换)应用于顶点的更新状态。以传统的图卷积网络(GCN)为例,在每一层中,一个顶点使用求和函数聚合邻居状态和自己的状态,然后应用单层 MLP 来转换新状态。如果层数为 𝐿,则此类过程重复 𝐿次。第𝐿层生成的顶点状态被下游任务使用,如节点分类、链路预测等。在过去的几年中,许多研究工作在图神经网络模型的设计方面取得了显着进展。突出的模型包括 GCN [108]、GraphSAGE [46]、GAT [95]、GIN [118] 和许多其他特定于应用程序的 GNN 模型 [128、130]。迄今为止,有数十项调查审查了 GNN 模型 [114、116、135、147]。另一方面,为了有效地开发不同的 GNN 模型,基于各种深度学习库 [8、14、34、44、68、104] 提出了许多面向 GNN 的框架。提出了许多新的优化来加速 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 训练。它将模型训练的全部工作量分配给一组workers,所有workers并行处理工作量。然而,由于 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 通信协议。该分类法不仅涵盖了mini-batch 分布式 GNN 训练和全图分布式 GNN 训练中使用的优化技术,还讨论了从图处理到模型执行的技术。我们仔细审查了每个类别中的现有技术,然后描述了来自工业界或学术界的 23 个具有代表性的分布式 GNN 系统和框架。最后,我们简要讨论分布式 GNN 训练的未来方向。
本次调查的贡献如下:
• 这是第一个关注高效分布式 GNN 训练优化技术的调查,并且它可以帮助研究人员快速了解分布式 GNN 训练的概况。
• 通过考虑端到端分布式GNN训练的生命周期,我们引入了分布式GNN训练技术的新分类。在高层,新的分类法由四个正交类别组成,即GNN数据分区、GNN批处理生成、GNN执行模型和GNN通信协议。
• 我们为新分类法中的每个类别提供了详细而全面的技术摘要。
• 我们审查了从工业界到学术界的 23 个具有代表性的分布式 GNN 培训系统和框架。
• 我们讨论了分布式GNN 训练的未来方向。
2 图神经网络的预测
图和图分区。图 𝐺 = (𝑉 , 𝐸) 由顶点集 𝑉 和边集 𝐸 组成。顶点𝑣 ∈ 𝑉 的邻居集用𝑁 (𝑣)表示,𝑣 的程度用𝑑 = |𝑁 (𝑣)|表示。在有向图中,顶点𝑣 ∈ 𝑉 的传入和传出邻居的集合分别用𝑁𝑖𝑛 (𝑣) 和𝑁𝑜𝑢𝑡 (𝑣) 表示,对应度用𝑑𝑖𝑛 = |𝑁𝑖𝑛 (𝑣)|和𝑑𝑜𝑢𝑡 = |𝑁𝑜𝑢𝑡 (𝑣)|.一条边 𝑒 ∈ 𝐸 由 (𝑢, 𝑣) 表示,其中 𝑢, 𝑣 ∈ 𝑉。𝐺的相邻矩阵由𝑨表示,其中当(𝑢,𝑣)∈𝐸时,𝑨的条目𝑎𝑢𝑣等于1,否则𝑎𝑢𝑣=0。图中的每个顶点𝑣都有一个初始特征向量𝒙𝒗 ∈ R𝐷 ,其中𝐷是特征向量的维数,图的特征矩阵由𝑿表示。
可以在分布式设置中对图形进行分区。对于边切割图分区𝐺 = (𝑉 , 𝐸) 分为𝑃 分区{𝐺𝑖 = (𝑉𝑖, 𝐸𝑖 )}, 1 ≤ 𝑖, 𝑗 ≤ 𝑃, 满足𝑉 = ∪𝑉𝑖, 𝑉𝑖 ∩ 𝑉𝑗 = ∅, 1 ≤ 𝑖, 𝑗 ≤ 𝑃. 交叉边的端点称为边界顶点𝐺 是内部顶点。对于顶点切割图分区𝑃 分区应满足𝐸 = ∪𝐸𝑖, 𝐸𝑖 ∩ 𝐸𝑗 = ∅, 1 ≤ 𝑖, 𝑗 ≤ 𝑃. 顶点可以在分区之间复制,并且顶点的复制数量𝑣 称为复制因子。当顶点的复制因子大于1时,该顶点称为边界顶点𝐺 是内部顶点。
图形神经网络(GNNs)。给定一个图表𝐺 具有相邻矩阵𝑨 和特征矩阵𝑿 其中每行是初始特征向量𝒙𝒗 顶点的𝑣 在图表中𝐺, 一𝑙-GNN中的第三层通过聚合来自相应邻域的特征来更新顶点特征,该过程可以在如下矩阵视图中形式化
其中,𝑯 𝒍 是隐藏嵌入矩阵,并且𝑯 0 = 𝑿 , 𝑾 𝒍−1是模型权重𝑨 是标准化的𝑨 和𝜎 是非线性函数,例如Relu、Sigmoid等。等式1是GNN计算的全局视图。GNN计算的局部视图是单个顶点的计算。给定一个顶点𝑣, 局部计算𝑙-消息传递模式[41]中的第层可以形式化如下
其中,⊕是聚合函数,𝜙是更新函数,𝒉𝒍-1𝒆𝒖,𝒗是第(𝑙-1)层的边(𝑢,𝑣)的隐藏嵌入。
GNN 训练方法。为了训练一个 GNN 模型,一个基本的解决方案是在单个 epoch 中更新一次模型参数,这种训练方法称为全图 GNN 训练。然而,全图 GNN 训练是内存密集型的 [108],因为它需要访问整个训练图并且不能扩展到大图。或者,我们可以选择在单个 epoch 的过程中多次更新 GNN 模型参数。这被称为小批量 GNN 训练。一个 epoch 被分成多次迭代,每次迭代用 batch(也就是整个训练数据集的一个子集)更新 GNN 模型。一个小批量可以在一台机器上处理,也可以由多台机器协同处理(称为大批量训练)。图形采样是一种为 GNN 模型生成小批量的流行工具。到目前为止,已经提出了许多基于采样的 GNN [17、23、46、56、125、151],根据不同类型,它们可以分为节点采样、层采样和子图采样。不同的批量生成方法会影响训练效率和训练精度。
3 分布式GNN训练与挑战
3.1 通用分布式GNN训练流水线
为了更好地理解端到端分布式 GNN 训练的一般工作流程,我们将训练流水线分为数据划分、GNN 模型优化和梯度聚合三个阶段。图 1 可视化了端到端分布式 GNN 训练工作流程的高级抽象。
数据分区。这是一个支持分布式训练的预处理阶段。它将输入数据(即图形和特征)分配给 一组工人。考虑到 GNN 中的训练数据是相互依赖的,数据划分阶段变得比传统分布式机器学习中的更为复杂。如图 1 所示,分区数据(即子图)之间的跨worker边暗示了数据依赖性。如果我们承认分区之间的数据依赖性,分布式训练效率会因通信而降低;如果我们简单地忽略数据依赖性,模型的准确性就会被破坏。因此,数据分区是端到端分布式 GNN 训练效率的关键阶段。
GNN 模型优化。这是分布式 GNN 训练的核心阶段,它执行 GNN 模型的训练逻辑(即前向计算和反向计算)。我们进一步将这个阶段分为计算图生成、执行模型和通信协议。首先,每个 worker 从分区的输入图和特征生成一个计算图,然后用计算图执行 GNN 模型以计算损失和梯度。然而,由于数据依赖性,计算图的生成和 GNN 模型的执行与传统的深度学习模型有很大不同。如果不访问远程输入数据,可能无法正确生成小批量训练策略的计算图。执行模型涉及 GNN 模型的层图聚合,并且聚合表现出不规则的数据访问模式。在分布式全图训练中,各层中的图聚合需要通过通信协议访问顶点远程邻居中的隐藏特征,并且还应考虑层与层之间的同步模式。因此,与传统的分布式机器学习相比,由于 GNN 中的数据依赖性,计算图生成和模型执行变得更加复杂。
梯度聚合。该阶段负责聚合最新的局部梯度以获得全局梯度并更新模型参数。在 GNN 的背景下,模型尺寸通常很小,模型更新逻辑与其他机器学习模型相同。经典分布式机器学习中现有的梯度聚合技术可以直接应用于分布式 GNN 训练。
3.2 分布式 GNN 训练的挑战
由于数据依赖性,有效训练分布式 GNN 并非易事。我们总结了以下三大挑战。
挑战 #1:海量特征交流。对于分布式 GNN 训练,在 GNN 模型优化阶段会产生大量特征通信。当GNN 模型以 mini-batch 方式训练时,计算图生成(即 mini-batch 构造)需要访问远程图和特征,如图 1 中的操作 1○ 以创建用于本地训练的批次。即使在使用集中式图存储(例如 CPU 内存)进行的多 GPU 训练中,计算图的生成也会导致从图存储到worker(例如 GPU 内存)的大量数据移动。现有的实证研究表明,小批量构建的成本成为端到端训练中的瓶颈 [57, 140]。当 GNN 模型以全图方式训练时,计算图可以与分区图相同而无需通信,但是每层中的图聚合需要访问顶点远程邻居中的隐藏特征(在图 1 中的操作 2 ○ )导致大量隐藏特征(或嵌入)通信。总之,GNN 模型的分布式小批量训练和分布式全图训练都受到海量特征通信的影响。
挑战 #2:模型准确性的损失。 Mini-batch GNN 训练比全图训练更具可扩展性,因此被大多数现有的分布式 GNN 训练系统采用。然而,随着模型深度的增加,精确的小批量训练会遇到邻居爆炸问题。事实上的解决方案是通过采样或忽略跨worker边缘来构造一个近似的小批量。虽然近似的 mini-batches 提高了训练效率,但它不能保证理论上模型的收敛性 [58]。因此,对于分布式 mini-batch GNN 训练,我们需要在模型精度和训练效率之间做出权衡。此外,有发展全图分布式GNN训练的趋势,这种训练具有收敛性保证,能够达到与单节点训练相同的模型精度。
挑战 #3:工作负载不平衡。工作负载平衡是分布式计算中的一个内在问题。然而,GNN 模型的各种工作负载特征增加了在worker之间划分训练工作负载平衡的难度。因为很难以简单统一的方式对 GNN 工作负载进行建模。如果没有正式的成本模型,经典的图划分算法就不能用来平衡worker之间的 GNN 工作量。此外,分布式 mini-batch GNN 训练要求每个 worker 处理具有相同批大小(即子图大小)的相同数量的 mini-batch,而不是简单地平衡子图中的顶点数量。综上所述,在分布式环境下训练 GNN 很容易遇到工作负载不平衡的情况,从而导致 worker 相互等待,破坏训练效率。
3.3 分布式 GNN 训练技术的分类
为了通过解决上述挑战实现分布式 GNN 训练并优化效率,过去几年提出了许多新技术。大多数现有作品从系统或框架的角度展示了自己的技术贡献,并没有很好地从技术角度回顾相关工作。在本次调查中,我们根据端到端分布式训练管道中的阶段组织分布式 GNN 特定技术,从而引入了一种新的分类法。通过这样的设计,我们将分布式 GNN 训练流水线中相同阶段优化的类似技术整理在一起,帮助读者充分理解分布式 GNN 训练中不同阶段的现有解决方案。
根据以往的实证研究,由于数据依赖性,分布式GNN训练的瓶颈一般出现在pipeline中的数据划分和GNN模型优化阶段,而新的 GNN 特定技术位于这些阶段。此外,各种训练策略(例如,小批量训练、全图训练)带来不同的工作负载模式,并导致在模型优化中使用不同的优化技术。例如,模型优化阶段的计算图生成对小批量训练很重要,而通信协议对全图训练很重要。因此,我们的新分类法将分布式 GNN 特定技术分为四类(即 GNN 数据分区、GNN 批处理生成、GNN 执行模型和 GNN 通信协议),如图 2 所示。在下文中,我们介绍了每个类别的概述类别。
GNN 数据分区。在这个类别中,我们回顾了分布式 GNN 训练的数据分区技术。数据分区的目标是平衡工作负载并最小化 GNN 工作负载的通信成本。 GNN训练是分布式图计算的一个实例,很多传统的图划分方法可以直接使用。然而,由于 GNN 工作负载的新特征,它们对于分布式 GNN 训练来说并不是最优的。研究人员付出了很多努力来设计指导传统图形划分方法的 GNN 友好成本模型。此外,图和特征是 GNN 中两种典型的数据类型,它们都是分区的。一些作品将特征与图结构解耦并独立划分。在第 4 节中,我们详细阐述了现有的 GNN 数据分区技术。
GNN 批量生成。在这个类别中,我们回顾了用于小批量分布式 GNN 训练的 GNN 批量生成技术。小批量生成的方法会影响训练效率和模型精度。图采样是为大规模 GNN 训练生成小批量的一种流行方法。然而,标准的图采样技术没有考虑分布式环境的因素,并且一个worker上的每个采样器都会频繁访问其他worker的数据,从而导致大量通信。最近,引入了几种针对分布式设置优化的新GNN 批量生成方法。我们进一步将它们分为分布式采样小批量生成和基于分区的小批量生成。此外,已经对缓存进行了广泛研究,以减少 GNN 批处理生成期间的通信。在第 5 节中,我们详细阐述了现有的 GNN 批量生成技术。
GNN 执行模型。在这个类别中,我们回顾了计算图生成和计算图执行的执行模型。在计算图的生成过程中,采样和特征提取是占总训练时间的两个主要操作。为了提高效率,提出了不同的执行模型以充分利用计算资源。图聚合应该在计算图生成之后进行,我们称之为计算图执行。在此阶段,每个节点的邻居状态在前向计算中聚合,梯度在后向计算中分散回邻居。由于数据依赖性和不规则的计算模式,传统的机器学习并行模型(例如,数据并行,模型并行等)对于图聚合不是最优的,特别是当特征向量是高维时。从图算子执行顺序的角度,我们将计算图执行模型分为one-short执行模型和chunk-based执行模型。从图聚合的同步时效性角度,我们将计算图执行模型分为同步执行模型和异步执行模型。基于去除不同的同步点,异步执行模型可以进一步分为两类。这些技术的详细描述在第 6 节中介绍。
GNN 通信协议。在这个类别中,我们回顾了分布式全图并行训练的通信协议。分布式全图训练的图聚合需要访问远程隐藏嵌入访问,因为顶点的邻居不能总是在图分区的本地。基于图聚合中的同步和异步执行模型,这些通信协议也可以分别分为同步通信协议和异步通信协议。前者假设在每一层中,模型访问最新的嵌入,而后者假设在每一层中模型被允许访问过时的嵌入(或历史嵌入)。第 7 节详细介绍了各种通信协议。
4 GNN 数据分区
在本节中,我们回顾了分布式 GNN 训练中 GNN 数据分区的现有技术。图 3 描述了这些技术的概述。考虑到图和特征是 GNN 中两种典型的数据类型,我们将划分方法分为图划分和特征划分。优化目标是工作负载平衡以及通信和计算最小化,旨在解决挑战 #1 和 #3。此外,成本模型是捕获 GNN 工作负载特征的另一个关键组件。在下文中,我们首先介绍成本模型,然后分别讨论图形分区和特征分区。
4.1 GNN的成本模型
成本模型对数据分区至关重要,因为它对 GNN 工作负载的计算和通信成本进行建模。通常,对于图形分析任务,我们使用顶点数来估计计算成本,使用交叉边数来估计通信成本 [63]。然而,这种简单的方法并不适用于 GNN 任务,因为它们不仅受顶点和交叉边数的影响,还受特征维度、层数和训练顶点分布的影响。研究人员提出了几种特定于 GNN 的成本模型,包括启发式模型、基于学习的模型和基于算子的模型。
启发式模型选择几个图形指标,通过简单的用户定义函数来估计成本。在流式图分区 [89、119]的上下文中引入了几种用于 GNN 工作负载的启发式模型,它将一个顶点或块一个一个地分配给一个分区。这样的模型为每个顶点或块定义了一个亲和力分数,该分数帮助顶点或块选择合适的分区。
假设 GNN 任务由 𝐺𝑁 (𝐿, 𝐺, 𝑉𝑡𝑟𝑎𝑖𝑛 , 𝑉𝑣𝑎𝑙𝑖𝑑 , 𝑉𝑡𝑒𝑠𝑡 , 𝑉𝑡𝑒𝑠𝑡 ) 表示,其中 𝐿 是层数,𝐺 是图,𝑉𝑡𝑟𝑎𝑖𝑛 、𝑉𝑣𝑎𝑙𝑖𝑑 和 𝑉𝑡𝑒𝑠𝑡 是训练、验证和测试的顶点集。将图𝐺划分为𝐾子图。在流图分区的上下文中,对于顶点或块的每次分配,令𝑃𝑖 (1 ≤ 𝑖 ≤ 𝐾) 为已分配给它的顶点集,𝑉 𝑖 𝑡𝑟𝑎𝑖𝑛 、𝑉 𝑖 𝑣𝑎𝑙𝑖𝑑 和𝑉 𝑖 𝑡𝑒𝑠𝑡 𝑡𝑒𝑠𝑡 𝑡𝑒𝑠𝑡 是对应的顶点集属于分区𝑃𝑖。
Lin 等人 [67] 为每个训练顶点 𝑣𝑡 ∈ 𝑉𝑡𝑟𝑎𝑖𝑛 定义一个亲和分数向量,其中每个分数代表顶点与分区的亲和力。令𝐼(𝑣𝑡)为训练顶点𝑣𝑡的neigh-hop in-neighbor集,𝑣𝑡相对于分区𝑃𝑖的得分定义如下,
其中,𝑉 𝑎𝑣𝑔 𝑡𝑟𝑎𝑖𝑛 = 𝑉𝑡𝑟𝑎𝑖𝑛𝐾 . 该得分函数隐式地平衡分区之间的训练顶点的数量。Liu等人[69]定义了关于块(或子图)的相似亲和性得分𝐵, 正式的定义是
其中,𝑃𝑎𝑣𝑔 = |𝑉 |𝐾 . Zheng等人[140]定义了块的亲和性得分𝐵 通过考虑所有训练、验证和测试顶点,公式如下
其中,𝐶𝑟𝑜𝑠𝑠𝐸𝑑𝑔𝑒 (𝑃𝑖, 𝐵)) 是中间交叉边的数量𝐵 和𝑃𝑖 , 𝛼, 𝛽, 𝛾 是用户手动设置的超参数。
基于学习的模型利用机器学习技术对GNN工作负载的复杂成本进行建模。其基本思想是通过特征工程手动提取特征,并应用经典机器学习来训练成本模型。基于学习的模型不仅能够使用静态图结构信息,还能够使用GNN工作负载的运行时统计信息来估计成本,从而实现比启发式模型更准确的估计。Jia等人[58]介绍了GNN计算成本估算的线性回归模型。该模型估计了单个GNN层的计算成本𝑙 关于任何输入图𝐺 = (𝑉 , 𝐸). 对于图中的每个顶点,他们选择五个特征(列于表1中),包括三个图结构特征和两个运行时特征。估计模型形式化如下,
其中,𝑤𝑖 (𝑙) 是层l的可训练参数,𝑥𝑖 (𝑣) 是v的第i个特征,以及𝑥𝑖 (𝐺) 总结了G中所有顶点的第i个特征。
Wang等人[102]使用多项式函数𝑓 以在一组手动选择的特征上估计顶点的计算成本。正式定义为
其中,𝑇 是由GNN模型定义的邻居类型(例如元路径[102])的数量,𝑛𝑖 是𝑖-th型,𝑚𝑖 是𝑖-第种类型的邻居实例(即𝑖-th类型具有𝑛 顶点,每个顶点都有特征尺寸𝑓 , 然后𝑚𝑖 = 𝑛 × 𝑓 ). 有关函数的详细示例,请参阅原著[102]𝑓 . 子图的总计算成本是子图中顶点的估计成本之和。
基于运算符的模型枚举GNN工作负载中的运营商,并通过对每个运营商的成本求和来估计总计算成本。Zhao等人[136]将GNN工作量的计算分为正向计算和反向计算。GNN层的正向计算分为聚合、线性变换和激活函数;而GNN层的反向计算分为朝向损失函数的梯度计算、嵌入梯度计算和梯度乘法。因此,计算嵌入的成本ℎ𝑙 𝑣 顶点的𝑣 在层中𝑙 前向和后向传播由𝑐𝑓 (𝑣, 𝑙) 和𝑐𝑏 (𝑣, 𝑙), 分别地
其中,𝑑𝑙 是𝑙-第GNN层|𝑁𝑣 | 是顶点的邻居数𝑣, 𝛼, 𝛽, 𝛾, 和𝜂 是可以通过测试实际运行时间来学习的恒定因素。最后,小批量的计算成本𝐵 通过将𝐵 从层1到层𝐿 如下所述,
其中,𝑁 𝑙 𝑢 表示图G中距离顶点u l跳的顶点。
4.2 GNN 中的图划分
GNN 是一种图计算任务。经典的图划分方法可以直接应用于 GNN 工作负载。许多分布式 GNN 训练算法和系统,如 AliGraph [150]、DistGNN [74] 采用 METIS [63]、Vertex-cut [43]、Edge-cut [63]、流图分区 [89] 和其他图分区方法 [ 10] 在不同的应用中训练 GNN 模型。最近,Tripathy 等人。[94] 对 GCN 模型的 1-D、2-D、1.5-D 和 3-D 图划分方法进行了实证研究。然而,由于 GNN 工作负载的独特性,经典的图分区方法并不是平衡 GNN 工作负载同时最小化通信成本的最佳方法。
在分布式小批量GNN训练中,为了平衡工作负载,我们需要平衡worker之间的训练顶点数量(即每个epoch的批数)。因此,引入了新的优化目标。Zheng等人[142]将DistDGL中的图划分问题表述为多约束划分问题,旨在平衡每个分区中的训练/验证/测试顶点/边。他们采用METIS中的多约束机制来实现同构图和异构图的定制图划分目标[143],并将该机制进一步扩展为两级划分策略,以处理多GPU和多机器的情况,并平衡计算。Lin等人[67]应用流式图分区–(LGD)[89],并使用新的亲和度分数(E.q.3)来平衡GNN工作负载。此外,它们确保分区中的每个顶点都有其完整的𝐿-跳邻居以避免采样期间的大规模通信。最近,基于块的流图分割方法[69,140]被引入用于小批量分布式GNN训练。他们首先使用多源BFS将图划分为许多小块,然后应用具有自定义亲和度分数(例如4、5)的贪婪分配启发式算法来划分粗化的图,并通过将块映射回原始图中的顶点来解开基于块的图划分。除了工作负载平衡,一些工作还可以最小化总计算成本。Zhao等人[136]证明了计算成本函数(E.q.11)是一个亚模函数,并使用METIS将图划分为内聚的小批量,达到2−2/𝑀 最佳计算成本的近似值。
在分布式全图GNN训练中,我们需要在一组worker之间平衡每个GNN层的工作量,同时最小化worker之间的嵌入通信。为了从不同的GNN模型中建模复杂成本,应用了基于学习的成本模型。在模型感知成本模型之上,采用了现有的图划分方法。例如,Jia等人[58]利用基于线性回归的成本模型(E.q.7),并应用范围图划分来平衡成本模型估计的工作量。此外,范围图分区保证每个分区都包含连续编号的顶点,从而降低了CPU和GPU之间的数据移动成本。Wang等人[102]使用E.q.8来估计分区的计算成本,并应用应用程序驱动的图分区方法[31]来生成工作负载平衡计划,该计划自适应地减轻工作负载并最小化通信成本。基于顶点切割分区,Hoang等人[48]利用2D笛卡尔顶点切割来提高可扩展性。
4.3 GNN 中的特征划分
作为 GNN 中重要的数据类型,特征也需要进行划分。大多数分布式 GNN 训练解决方案将特征视为图的顶点属性,并且它们与图分区一起进行分区。例如,如果用边切法对图进行分区,那么每个顶点特征都存储在对应顶点所在的分区中。换句话说,特征矩阵是按行划分的。
考虑到与图结构相比特征的不同处理模式,一行工作独立于图划分特征矩阵。当通过 2D 分区方法对图进行分区时,Tripathy 等人 [94] 也使用 2D 分区方法对特征矩阵进行分区,而 Vasimuddin等人 [74]使用逐行方法对具有复制的特征矩阵进行分区,并确保每个顶点在局部具有完整的特征。甘地等[40]深入分析分布式 mini-batch GNN 训练中的通信模式,发现经典的图划分方法仅降低了第一跳过程的通信成本。此外,复杂的图分区(例如 METIS)是一个昂贵的预处理步骤。他们分别介绍了对图和特征进行分区。该图通过随机分区进行分区,以避免昂贵的预处理。输入特征沿特征维度进行分区(也称为按列分区)。换句话说,每个分区包含所有顶点特征的子列。利用这种数据分区策略,他们进一步设计了第 6.2.4 节中介绍的新执行模型,并如第 7.2.2 节所述以最少的通信训练 GNN 模型,尤其是当输入特征具有高维且隐藏特征具有低维时方面。同样,Wolf 等人 [109] 还使用列式特征划分方法来分布式训练超宽 GCN 模型。不同的是,在这种方法中,图是用 METIS 划分的,输入和中间特征都沿着特征维度划分。
5 GNN 批量生成
小批量 GNN 训练是将 GNN 扩展到大图的常用方法。图抽样是在独立模式下生成小批量的实际工具。到目前为止,已经提出了许多基于采样的 GNN [17、23、46、56、125、151],它们根据不同类型的抽样方法进行抽样,可以分为顶点采样、层采样和子图采样。不同的批量生成方法会影响训练效率和训练精度。为了避免图采样成为瓶颈,对高效的 GNN 数据加载器进行了一些探索 [5、25、27、76]。
在 mini-batch 分布式 GNN 训练中,数据依赖为批量生成过程带来大量通信。为了提高分布式设置中的训练效率,提出了几种特定于分布式训练的新 GNN 批量生成技术,并解决了挑战 #1 和挑战 #2。如图 4 所示,一种解决方案是通过分布式采样生成 mini-batch,另一种是直接使用局部分区(或子图)作为 mini-batch。下面,我们将描述这些技术。
5.1 分布式采样的小批量生成
基于现有的基于采样的 GNN,通过实现分布式采样可以直接获得分布式版本。换句话说,我们使用分布式图采样在大图上创建小批量。大多数现有的分布式 GNN 系统,如 AliGraph [150]、DistDGL [142]、BGL [69],都遵循这一思想。然而,在分布式环境中,worker 上的基本采样器将频繁访问其他worker 的数据,从而导致大量通信。在多 GPU 设置中,虽然图形和特征可能集中存储在 CPU 内存中,但每个 GPU 都需要通过大量 CPU-GPU 数据移动从 CPU 内存访问 mini-batch。
为了减少分布式计算中的通信,缓存是改善数据局部性的常用优化机制。在系统层面,已经提出了许多面向GNN的缓存策略。缓存的基本思想是在本地存储对远程顶点的频繁访问。Zheng等人[142]建议为DistDGL中的每个分区图复制局部边界顶点的远程邻居,并确保局部的每个顶点都具有完整的邻居,以便小批生成器(例如,采样器)可以在不与其他分区通信的情况下进行局部计算。然而,这种方法仅减少了由直接(即,一跳)邻居接入引起的通信,并且不知道接入频率。Zhu等人[150]引入了一种称为𝑣的第𝑙重要性, 表示为𝐼𝑚𝑝 (𝑙) (𝑣). 度量正式定义为𝐼𝑚𝑝 (𝑙) (𝑣) = 𝐷 (𝑙 )𝑖 (𝑣)𝐷 (𝑙 )𝑜 (𝑣) , 其中,𝐷 (𝑙)𝑖 (𝑣) 和𝐷 (𝑙)𝑜 (𝑣) 是𝑙-跳入和跳出顶点的邻居𝑣。Li等人[67]设计了一个静态GPU缓存,用于存储多GPU训练的高出度顶点。静态缓存可以减少CPU和GPU之间的功能移动。Min等人[75]引入了加权反向pagerank(加权R-pagerank)[6],将顶点的标记状态结合到反向pagerank方法中,以识别频繁访问顶点。得分较高的热顶点特征被缓存在GPU内存中,其余冷顶点特征被放置在CPU内存中。此外,由于NVLink和PCIe之间的带宽不同,他们将热顶点功能分为最热的和次热的。最热门的顶点特征在多个GPU上复制。下一个最热的数据分散在多个GPU上,可以从对等GPU内存访问。
上述静态缓存策略不能适应各种图形采样工作负载。Liu等人[69]提出了一种动态缓存机制,通过利用顶点特征的分布来应用FIFO策略。为了提高缓存命中率,他们为小批量生成引入了邻近感知排序。该顺序以BFS顺序约束训练顶点序列。为了确保模型收敛,他们通过以循环方式选择BFS序列并随机移动BFS序列,将随机性引入到基于BFS的序列中。Yang等人[122]提出了一种更稳健的缓存方法,称为基于预采样的缓存策略。给定图G、采样算法a和训练集T,它首先执行𝐾 采样阶段收集统计信息,然后使用统计信息识别数据的频率(或热度),最后缓存具有更高热度的数据。基于预采样的缓存策略考虑了影响数据访问频率的所有三个因素,并且对于各种GNN工作负载,它比以前的缓存策略更稳健。
除了缓存策略,另一种方法是设计新的分布式采样技术,这些技术具有通信效率并保持模型准确性。通信高效采样的一个基本思想是先对局部顶点进行采样,但这会为生成的小批量引入偏差。受线性加权采样方法 [17、151] 的启发,Jiang 等人[60] 提出了一种用于邻居选择的倾斜线性加权采样。具体而言,倾斜采样通过因子s > 1 缩放局部顶点的采样权重,Jiang 等人从理论上讲,训练可以通过正确选择s的值来实现与线性加权采样相同的收敛速率。
5.2基于分区的小批量生成
根据图 1 中分布式 GNN 训练流水线的抽象,我们知道首先对图进行了分区。如果我们限制每个worker只能使用本地分区训练 GNN 模型,则可以避免大量通信。在这样的训练策略中,一个partition就是一个mini-batch,我们称之为partition-based mini-batch。 PSGD-PA [82] 是上述想法与参数服务器的直接实现。在 GraphTheta [65] 中,分区是通过社区检测算法获得的。
然而,忽略分区之间的交叉边会导致模型精度的损失。为了提高模型精度,引入了子图(即分区)展开。它通过将远程顶点复制到本地来保留边界顶点的本地结构信息。薛等人[120]使用METIS获得一组未重叠的子图,然后通过添加不属于子图的顶点的一跳邻居来扩展每个子图。类似地,Anged等人[3]使用广度优先方法来扩展𝑙-跳子图。Zhao等人[138]定义了基于蒙特卡洛的顶点重要性测量,并使用深度优先采样来扩展子图。他们进一步引入了基于方差的子图重要性和加权全局一致性,以减少子图方差的影响并提高模型精度。此外,Ramezani等人[82]引入了LCCG,以提高基于分区的小批量训练的模型精度。LCCG遵循参数服务器架构,并定期在参数服务器上执行全局校正操作以更新全局模型。实证结果表明,全局校正操作可以显著提高模型精度。
6 GNN 执行模型
图 1 中的 GNN 模型优化阶段由计算图生成和计算图执行两个子阶段组成。执行模型负责调度子阶段的操作,以实现高训练效率。这两个子阶段针对不同的训练范式具有不同的操作符。对于分布式小批量训练,采样和特征提取操作决定了训练效率,并使计算图生成变得昂贵。而对于分布式全图训练,由于worker之间的数据依赖,计算图的执行很复杂。
下面分别讨论计算图生成和计算图执行的执行模型。
6.1 计算图生成执行模型
采样和特征提取是计算图生成中的两个独特操作。它们是分布式 mini-batch GNN 训练的瓶颈。如[40, 69] 所示,采样和特征提取几乎占据了整个训练时间的 83% - 99%。低效率是由于大量的通信和资源竞争造成的。因此,提出了几种用于计算图生成的执行模型。
传统的执行模型依次对每个计算图(即batch)进行采样和特征提取。当应用基于CPU的图采样时,采样在CPU上执行以构建图结构,针对图的特征提取在具有缓存支持的GPU上执行,训练也在GPU上执行。当应用基于 GPU 的采样 [57、79]时,管道中的所有操作都在单个设备上执行。这种传统的执行模型在现有的 GNN 训练系统中很流行[67、142、150]。传统执行模型上下文中的优化在于采样操作的优化和特征提取的缓存策略,如第 4节和第 5 节所述。但是,传统执行模型带来了资源竞争。例如,在应用基于CPU的图采样时,处理图采样请求和子图构造会竞争CPU;当使用基于 GPU 的图形采样时,图形采样和特征提取会竞争 GPU 的内存。
分解执行模型使用专用设备或隔离资源来执行计算图生成中的单个操作。实证研究表明,在mini-batch GNN 训练期间,一个操作在不同时期共享大量数据并保留期间数据局部性。这一观察结果促使分解执行模型避免资源争用并改善数据局部性。杨等[122] 介绍了一种分解执行模型,可以消除 GPU 内存上的资源争用。新模型为各个设备分配采样操作员和训练操作员。更具体地说,它使用一些专用的 GPU 进行图形采样,其余的 GPU 用于特征提取和模型训练。在分解设计的帮助下,新的执行模型能够对更大的图进行采样,而用于特征提取的 GPU 可以缓存更多特征。在使用基于 CPU 的图形采样时提高 CPU 利用率。刘等人 [69] 将隔离的资源(即 CPU)分配给不同的操作,并通过调整资源分配来平衡每个操作的执行时间。他们引入了一种基于分析的资源分配方法,以提高 CPU 利用率和处理效率。新的资源分配方法首先描述了每个操作员的成本,并解决了一个优化问题,在资源约束下最小化所有操作的最大完成时间。
算子并行执行模型通过混合管道中所有算子的执行来实现批间并行并并行生成多计算图。郑等[143] 将计算图生成分解为多个阶段(例如,图采样、邻域子图构造、特征提取),每个阶段都可以视为一个运算符。不同算子之间存在依赖关系,而 mini-batch 之间不存在依赖关系。因此,可以将不同mini-batch中的算子调度到一个pipeline中,不依赖的算子可以并行执行。这样,后一批的计算图生成可以在前一批计算图生成之后立即或同时进行,而不是等待前一批完成计算图执行阶段。张等[140] 采用类似的思想来解耦计算图的生成,同时将不同算子之间的依赖关系进一步建模到 DAG 中。将每个算子视为 DAG 中的一个节点,通过考虑 DAG 之间(粗粒度)和算子之间(细粒度)的调度,提出了一种两级调度策略。通过这种两级调度,可以平衡基于 CPU 的采样和基于 GPU 的特征提取的成本,以提高 CPU 利用率并减少端到端的 GNN 训练时间。
具有拉-推并行性的执行模型。以上所有执行模型都是先采样构造子图,然后根据子图提取特征。当输入特征是高维时,特征提取引起的通信成为端到端 GNN 训练管道中的瓶颈。 4.2 节中介绍的昂贵的高质量图分区是减少 Gandhi 等人通信的一种选择。 [40] 引入了一种新的拉推并行执行模型。新模型使用快速哈希分区对图进行分区,使用列式哈希分区对特征矩阵进行分区。它通过在本地拉远程顶点和边来构造子图,然后将构造的子图推送给所有的worker。然后每个worker并行地为所有接收到的子图提取部分特征。借助pull-push并行性,执行模型将昂贵的输入特征移动替换为光图结构移动,从而提高了计算图生成的效率。
6.2 计算图执行模型
生成计算图后,它们由分布式设置中的一组worker执行。对于分布式小批量 GNN 训练,并行化小批量计算图的执行是微不足道的,因为每个计算图都完全存储在单个 worker 中。对于分布式全图或大批量 GNN 训练,计算图的执行在一组工作人员之间进行分区。由于计算图中的数据依赖性,实现 GNN 训练的高效性并非易事。到目前为止,已经提出了许多针对GNN训练的计算图执行模型。