【调研】分布式GNN训练的综述汇报

转载请注明出处：小锋学长生活大爆炸[xfxuezhang.cn]

　　大家好，今天跟大家分享一下我最近看的两篇关于GNN的综述文章。整合了一下他俩的内容，相当于给大家简单介绍一下分布式GNN训练的全流程。
　　以下内容结合了我的理解，若有不对的地方，欢迎大家指出。
　　内容比较多，可能要占用大家30~60分钟的时间。

　　与其他神经网络相比，GNN可以有效地表示非欧式域。并且CNN可以看做是GNN的特例。它还可以与CNN、强化学习等技术组合使用。
　　GNN的应用场景非常广泛，包括系统建模、图像、文字、社交网络、图生成、电子设计等。
右图是GNN模型训练的典型架构图。左侧将图形数据作为输入，然后通过GNN模型学习图形中每个顶点的向量表示。最后输出的学习的表示，可用于下游任务，比如顶点预测、链接预测和图预测等。

　　这里简单统计了Arxiv和Google Scholar两个平台上，近几年标题中包含GNN或分布式GNN的文献数量。
　　可以看到，GNN自2018年兴起以来，相关的论文数量呈现出了指数级的增涨。
此外，分布式GNN训练的贡献在2019年开始显现，目前也呈现出了快速增长的趋势。这是因为工业界和学术界对缩短GNN模型训练时间的要求很高，因此越来越受到关注。
　　不过一些属于GNN，但标题中未包含关键词的论文没有统计进来。如果算上，预计会更多。

　　除了学术界，各大互联网企业也在积极部署自己的图计算平台，国内如腾讯、阿里、百度等，都已推出了相应的方案。
　　可见，图计算、图神经网络获得了各界一致的认可。

　　GNN模型由一个或多个层组成，而GNN模型的训练主要包括Aggregation邻居聚合和Combination后的神经网络操作。在聚合步骤中，聚合函数用于为每个目标顶点聚合来自前一个GNN层的传入相邻顶点的特征向量。在组合步骤中，组合函数使用神经网络操作变换每个顶点的聚合特征向量。
　　但随着大数据时代的发展，图的规模越来越大，甚至达到百亿个节点和边。训练这种GNN需要大量的计算和存储等资源。因此，分布式的GNN训练成了大图训练的有效解决方式。
　　根据分布式GNN的特点，论文从不同角度描述了现有工作对其的优化技术。通过合并两篇综述的主要内容，下面将从这四个方面简单介绍。

　　为了在多个计算节点上协作，需要进行数据分区。
　　但图数据是具有依赖性的，如果我们考虑分区之间的数据依赖性，分布式训练效率会因通信而降低；如果我们简单地忽略数据依赖性，模型的准确性就会被破坏。因此，数据分区是端到端分布式 GNN 训练效率的关键阶段。
　　在考虑依赖性的情况下，数据分区应以平衡工作负载和最小化通信成本为目标。

　　GNN任务由于受顶点数、交叉边数、特征维度、层数、顶点分布等因素的影响，简单的分区方法(如随机划分)，效果不是很好，因此学者们提出了包括基于启发式、基于学习和基于算子等几种特定于 GNN 的成本模型。
　　通过对 GNN 工作负载的计算和通信成本进行建模，可以更好地指导数据分区。

　　数据分区通常是按图划分，但由于现在的节点特征往往是高维向量或张量，因此大家也开始关注特征的划分，或者将两者结合。
　　图的划分可以分为Full-batch和mini-batch两种。Full-batch是将整张图输入训练，而mini-batch是通过采样生成小批量再输入训练。
　　Full-batch的挑战在于工作负载不平衡和大量通信。Mini-batch的挑战在于采样性能的不足。
　　后面会具体来讲。

　　特征的划分可以按列或按行来进行。
　　两种划分方式的示意图如右图所示。

　　首先介绍一下full-batch training。
　　由于每一轮都涉及整个原始图形数据，因此每一轮都需要大量的计算和大量的内存占用。为了解决这一问题，分布式全批训练主要采用工作负载划分方法：即分割图以生成小工作负载，并将其移交给不同的计算节点。
　　根据预处理阶段是否预设了工作负载，进一步分为基于调度工作负载的执行和基于预设工作负载的执行。

　　左边是基于调度工作负载的执行。
　　这里的Leader负责存储模型参数和图形数据，还负责生成和调度chunk，以及整合各worker的梯度并推进计算。而worker负责模型计算和梯度聚合。
　　我们看一下大致流程：
　　Leader首先分割Chunk，由于节点的邻居节点数量不同，因此chunk的大小也不一样。然后chunk被分发到各个worker。Worker根据需要从其他worker聚合邻域信息，并组合执行模型计算，然后将梯度发送回leader。Leader整合所有内容，并更新模型参数，将最终梯度推送给各个worker。Worker更新自己的参数。这就完成了一个epoch的训练。
　　对于这种方式，优化的方向包括以下几点。
　　首先是以顶点为中心的工作负载分区，常见的有1D、1.5D、2D、3D方式。如1D只需与左右方向的worker通信，2D只需与上下左右的worker通信，这可以提高计算节点的数据访问效率。
　　第二点，可以借助于成本模型来实现平衡工作负载的生成，比如线性回归成本模型。
　　传输优化可以减少数据传输的开销，比如可以通过缓存来减少重复数据的传输，或者通过允许计算组件或节点发送数据来合理的选择传输源，也就是说每个节点都可作为数据发送方。
　　最后还可以通过对特征的精细化拆分，来充分提高GPU线程级别的并行度。

　　右边基于预设工作负载的执行是在预处理阶段，按照一些策略，预先对数据进行了划分，以便考虑节点之间的相关性。
　　然后每个worker负责完成其子图中所有顶点的计算任务。
　　Worker根据需要从其他worker聚合邻域信息。并组合执行模型计算。然后将信息发送给参数服务器或中央节点。

之后由参数服务器或中央节点聚合梯度并向各个 worker 广播更新后的模型参数。

　　最后Worker收到后更新本地模型。
　　(这里的参数服务器或中央节点与前面的leader类似，但负责的工作较少，比如模型参数的管理和分发；而leader还需要负责整体的数据和计算的调度和分配)
　　不过这里的参数服务器或中央节点是可选的，因为模型参数是复制的。但这也意味着每一轮都需要进行梯度同步，以保证模型参数跨节点的一致性。

　　对于这种方式，优化的方向可以包括以下几点。
　　首先，预分区时保证子图的大小相似可以提高工作负载平衡，另外最小化边切割的数量可以减少聚合步骤中的通信开销。
　　传输优化与前面的类似，都是要减少通信开销。
　　延迟聚合允许worker在聚合步骤中使用旧的传输数据，这样可以将通信与计算重叠，从而减少通信开销。但为了保证收敛性和最终的准确性，延迟聚合主要基于有界异步形式，关键策略是限制最快的worker和最慢的worker之间的迭代次数在可接受范围。
　　激活再生成是通过在前向传播过程中存储所有激活到磁盘中，并在反向传播期间直接从磁盘重新计算或加载激活，从而减少内存压力。

　　分布式full-batch由于每一轮的计算都涉及到整个图形，虽然精度会更高，但通信量很大，并且内存容量需求高。
　　相比之下，分布式mini-batch每一轮的计算只涉及小批量，因此它触发的通信量更少，需要的内存容量也更少。
　　分布式mini-batch包括采样、模型计算和梯度同步这三个阶段。根据采样和模型计算是否解耦，可以分为基于单独样本的执行和基于联合样本的执行。

　　基于单个样本的执行涉及多个Sampler和Worker。
　　sampler首先对图数据进行采样生成mini-batch，然后将生成的mini-batch发送给workers。 worker 执行 mini-batch的计算并与其他 worker 进行梯度同步以更新模型参数。以此循环。
　　通过为sampler提供足够的计算资源来为worker准备小批量，可以在没有停顿的情况下执行计算。
　　并且由于mini-batch之间的数据量和计算量是相似的。而且GNN模型的尺寸较小，因此梯度同步的开销也比较小。因此，关键的优化点是如何加快采样过程，使其及时提供足够的小批量，避免停顿。
　　优化技术有以下几点。
　　首先，并行化mini-batch是指使用CPU的多线程来加快采样。
　　动态mini-batch分配是指不限制采样器和分配器的一一对应关系，还可以通过使用无锁输入队列来实现动态mini-batch的分配。
　　Mini-batch传输流水线将一个mini-batch的传输分成多个阶段，每个阶段对应着数据传输、计算和梯度更新等操作，这些操作可以并行执行，从而可以提高训练的效率。（具体还没看懂）
　　（对于 GNN 的聚合操作，需要为每个顶点聚合相邻顶点的特征。这意味着计算实际上是由边决定的。）
　　基于边划分的并行聚合主要思想是提升分区数据的独立性，从而提高聚合操作的并行性。比如在每个mini-batch 中，根据目标顶点将边划分为多个分区，以确保具有相同目标顶点的边可以位于同一分区中。然后这些分区由多个线程独立处理。由于这些线程之间没有数据冲突，因此聚合操作可以有效地并行完成。
　　右图中，在基于联合样本的执行中，图在预处理过程中通过分区被分割成子图。每个worker持有一个子图和模型参数的副本。worker对自己的子图进行采样以生成mini-batch。由于图是分区的，因此可能需要查询其他worker以获得目标顶点的邻居信息。然后对mini-batch进行模型计算并通过与其他worker的交流来同步更新模型。
　　对于这个模式，主要关心的是如何让worker的计算更加独立，从而获得更好的性能。
　　第一个优化点是通过局部感知技术，将图划分为局部性较好的子图，旨在通过减少节点之间的通信来使节点的计算更加独立。
　　重叠分区是指在对图进行分区时复制目标顶点的相邻顶点，从而减少甚至完全消除采样阶段计算节点之间的数据传输（仅限k-hop）。
　　细化的独立执行是让每个 worker 独立计算和更新模型参数，然后使用额外的 refine 操作定期对模型参数进行平均，从而进一步提升并行性。
然后是缓存技术，将来自其他worker的常用顶点信息进行缓存，以减少数据传输的开销。

　　这部分是关于通信协议，分为同步通信和异步通信。

　　将单个 GNN 层的前向计算分为四个算子: Scatter (SC)、ApplyEdge(AE)、Gather (GA) 和 AppleVertex (AV)。SC 和 GA 是两个图操作，其中顶点特征分别沿边scattered并gathered到目标顶点。AE 和 AV 可能包含神经网络 (NN) 操作，它们分别直接处理目标顶点的边缘特征或聚合特征。

　　在同步执行模型中，结合前向和反向传播，将执行流水线分为八个阶段。其中，有两个阶段涉及到边界顶点状态的传递，即GA和▽GA。
在▽GA中，边界顶点的梯度应该被发送回它们所属的worker。因此，GA和▽GA是同步执行模型中的两个同步点。在这些点上，必须阻塞执行流程，直到所有通信完成。

　　在同步执行模型中通信协议包括4种。
　　第一种是基于广播的协议。在 GNN 训练过程中，每个 worker 直接将自己的梯度或消息广播给所有的worker 进行同步，保证一个 worker 上的顶点具有完整的邻域。不过工作负载不平衡问题会破坏广播的效率，这会导致计算停滞。并且广播会产生很多冗余数据和网络拥塞。因此，这种协议适用于worker之间通信较少，且通信量不大的情况。
　　基于点对点的协议是一种在worker之间共享信息的细粒度通信方法，可以将梯度或消息直接发送给特定的worker，这种协议适用于worker之间通信频繁、通信量较大的情况。
　　基于流水线的通信协议，将计算和通信过程分成多个阶段，每个阶段完成一部分计算和通信，并将计算结果传递给下一个阶段。
　　对于共享内存通信。为了大规模训练大型GNN，可以利用CPU内存来检索所需的信息。因此可以将完整的图和特征存储在CPU共享内存中，而每个GPU的设备内存则被视为缓存。

　　为了降低通信成本并提高训练效率，提出了几种通信协议。

　　在执行流水线中移除计算图计算的同步点会引入 I 类异步。在 I类异步中，worker不等待最新信息的到达。相反，它使用从先前epoch中缓存或接收到的顶点的历史信息，来执行目标顶点的聚合。在▽GA阶段中也使用了类似的历史梯度。

在GNN训练期间，当权重参数需要更新时，会出现另一个同步点。通过删除这样的同步点，形成了II类异步。在II类异步中，如果一个worker完成了当前epoch的计算，它不一定要等待其他worker完成，而是全局更新权重以开始下一个epoch。相反，它可以使用由前几个epoch更新的历史权重并立即开始计算下一个epoch。

　　I类异步消除了聚合阶段的同步，II类异步消除了模型权重更新阶段的同步。 I 类异步专用于GNN 模型计算，而 II 类异步已广泛应用于传统深度学习，相应的技术可以直接引入 GNN 训练中。

　　异步 GNN 通信协议用于支持异步执行模型。通常，引入异步意味着在训练中应该使用陈旧的信息。
　　有三种流行的陈旧模型：固定Epoch的陈旧性、自适应Epoch的陈旧性和基于变化的陈旧性。
　　这几种模型都是为了使陈旧性有界，也就是前面提到过的保证延迟聚合是有界异步的。

　　这部分介绍了设计分布式GNN训练的软硬件平台。
　　软件平台给出的是目前已有的比较流行的框架。
　　硬件平台介绍了多CPU、多GPU以及混合设备的情况。
　　这些都比较好理解。

　　这是文献1中给出的不同软件框架用到的技术分类，供大家参考。

　　这是文献2中给出的不同软件框架用到的技术分类。

　　这是看到的另一篇综述，里面也提到了一些相关的。

　　再完整的回顾和总结一下。
　　对于一张大图、一个GNN模型、以及多台计算节点，如何去实现分布式并行加速。
　　常见的分布式DNN训练包括数据并行、模型并行和混合并行。对于分布式GNN来说，也可以从这几个方向出发。
　　但GNN模型一般很简单，甚至可能就两层，所以模型并行不是很有用。更多是从数据并行上来扩展，毕竟图数据可以是很大的。
　　对于计算节点，可以是CPU、GPU、Serverless或FPGA等的组合。

　　因此一轮分布式GNN的流程就变成了：大图分割为小图，将小图分发到各个节点进行模型计算，然后将各节点的梯度组合进行梯度更新。
　　这三个阶段中，每个阶段都有学者进行了优化研究，并考虑了“内存、通信、效率”等指标，或者说“内存需求高、负载不平衡、大量数据传输”等问题。

　　分割图时需要考虑如何更好的对图分区，以减少之后的通信，并如何提高batch生成和分发的效率。
　　对于大图，可以通过分区拆分成子图，或者采样生成mini-batch。
　　可以从图或特征的角度，按节点、分层、子图来划分。
　　还可以使用成本模型来指导更好的分区，或者在特征维度上拆分来探索特征级的并行性。为了减少数据传输的开销，可以使用缓存或者传输流水线来进行优化。

　　模型计算涉及到聚合和组合两个操作。聚合阶段需要从其他计算节点获取邻居信息，因此对于负载平衡和通信的要求较高。
　　可以使用缓存技术来优化，比如缓存重复数据，或缓存k-hop的邻居信息。对于传输路径和策略也可以进行优化。还可以设计如何让计算节点更加独立，以便充分发挥并行性。
　　而由于GNN模型比较简单，因此组合阶段一般没有什么优化点。

　　对于梯度的同步，除了传统的同步等待外，还可以允许使用k轮的旧数据，也可以让节点先独立计算，最后只是再做一个平均运算。

　　其他还有不同平台之间的传输加速优化方向可以研究。一般而言，CPU更适合用于采样，而GPU则用于计算。

这里是一些延伸阅读，供大家参考。

https://github.com/chwan1016/awesome-gnn-systems
https://github.com/zhenyuyangmq/awesome-graph-level-learning
https://lgd.gd/2020/04/30/分布式深度学习研究笔记/
https://www.cnblogs.com/rossiXYZ/p/14856464.html
https://github.com/dglai/WWW20-Hands-on-Tutorial/blob/master/_legacy/basic_apps/BasicTasks_pytorch.ipynb
https://towardsdatascience.com/graph-ml-in-2022-where-are-we-now-f7f8242599e0
Recommend：