论文精讲 | MindSpore Graph Learning设计原型——图神经网络框架Seastar的编译优化

论文标题

Seastar: Vertex-Centric Programming for Graph Neural Networks

论文链接

http://www.cse.cuhk.edu.hk/~jcheng/papers/seastar_eurosys21.pdf

Seastar发表于EuroSys 2021是香港中文大学James Cheng教授团队和华为合作出品，也是MindSpore Graph Learning的设计原型：

下面给大家解读一下其中的设计思路：

01

图神经网络框架

图神经网络被设计用于从社交网络、知识图谱等图结构数据中提取信息表征。从2005第一次被提出至今，图神经网络被广泛用于商品推荐，金融风控、计算机视觉、生物计算甚至材料科学等领域。不同于图像和文本数据，图结构数据具有不规则的特性，基于Tensor的深度学习框架不能很好地直接表达图操作和图神经网络。为了支持用户更好地表达图操作和图神经网络，各大科技公司纷纷提出自家的图神经网络框架例如DGL、PYG、Graph-Learn以及PGL等。

从表达的角度来说，很多框架的模型表达可以总结为针对GNN的算子库。部分框架提供了图的抽象，用户将特征向量转为Tensor存为图的属性，图提供接口执行图神经网络，并根据消息传播机制模型提供了消息和聚合两个函数，接收图的属性数据作为输入和输出。一些框架根据消息传播机制或SAGA(Scatter-Apply-Gather-Apply)模型提供计算模型不同阶段的算子，要求用户自己组织相应阶段算子的输入Tensor。部分框架偏重于GNN训练中的图采样过程，提供了采样操作的接口，模型计算则基于深度学习框架的算子实现。虽然消息传播机制等编程模型足以表达图神经网络，但是用户需要将图操作转化为整图Tensor计算代码逻辑。当需要实现复杂图操作时，这部分工作并不轻松。此外用户需要学习底层相关Tensor算子以及领域相关的API。

从运行的角度来说，很多框架在算子的泛化性和性能优化之间做tradeoff。一些框架倾向于泛化性，复用深度学习框架的Tensor计算优化，中间插入Scatter/Gather等算子将图数据转成Tensor，这样导致大量内存占用和数据读写，性能并不好。相反，一些框架为典型的算子提供了融合算子，性能提升不少，但是针对用户自定义算子无法进行优化，泛化性差。

有没有可能图神经网络的表达不再依赖Tensor独立成体系？在性能优化方面，框架能兼顾泛化性和优化效率吗？这是Seastar尝试做到的方向。Seastar为此提出了以点为中心的编程模型，支持Python风格的图神经网络表达，用户仅需用熟悉的Python语法实现中心节点的逻辑。针对GNN的运行模式，Seastar设计了一个通用的算子生成器，为GNN生成前向和反向都高效的融合算子。

下面将分别介绍Seastar的编程模型，算子生成器以及性能优化方法。

02

以点为中心的编程模型

从图神经网络的公式可以观察到，它主要描述了中心节点的特征通过邻居节点特征汇聚进行更新。为了支持用户根据公式进行以点为中心的编程，Seastar引入了一个函数装饰器，被装饰函数以节点作为输入，在被装饰函数内，用户可以利用Python语法实现节点的计算逻辑。

如下图基于以点为中心编程模型实现经典GAT网络，用户定义一个函数以节点v作为入参，在函数内用户通过v.innbs()获取邻居节点列表，遍历每个邻居节点u，获取节点特征，计算邻居节点与中心节点的特征交互得到邻边上的权重列表，然后将邻边权重与邻居节点进行加权平均，返回更新的中心节点特征。可以看到，以点为中心的编程实现的GNN代码与公式一一对应，用户可以轻松地实现GNN模型。

代码生成

不同于其他框架逐一执行算子并具体化算子的中间输出，可以很直接的与深度学习框架融合。Seastar需要将以点为中心的编程代码转换为基于Tensor算子的计算图并且与深度学习框架后端的运行时相结合。Seastar整体框架如下图，首先，将以点为中心的编程代码通过tracer记录下来并将之转换为前向计算图，同时进行用户编码正确性校验。具体的转换过程分两步：

第一步将前端以点为中心的代码中的算子记录下来转换为计算图DAG，为每个节点或边属性特征创建Tensor继承节点或边特征的数据类型、shape等作为算子的输入输出数据, 所有的深度学习算子及其实现用monky-patch的方式插入计算图。这里的计算图DAG已经可以基于深度学习框架编译执行，但是没有进行并行针对单个节点执行效率较低。

第二步则是将计算图DAG中Tensor以及算子与图结构或操作信息结合，为每个Tensor打上S(源节点)、D(目标节点)、E(边)等类型标签，算子则增加了A(邻居汇聚)类型，生成图感知的中间表达GIR。这些类型标签表明了全量数据的维度为节点粒度或边粒度，也表明了算子处理的数据类型，可以进行编码正确性校验，同时也为反向计算和优化提供了支持。

然后基于前向计算图GIR，根据链式求导法则生成反向计算图GIR，图类型标签用于校验反向逻辑的正确性。最后，根据GIR进行编译生成执行单元。Seastar将计算图根据是否进行融合优化分为融合单元和非融合单元，融合单元会基于Seastar融合模板进行融合编译，作为一个用户自定义算子插入后端，非融合单元直接调用后端实现。

编译优化

图神经网络的计算图GIR存在共有的S-E-A的计算模式，即从处理源节点粒度的特征(S)开始，然后再根据ID获取源节点、目的节点或边粒度特征进行边上的计算(E)，最后将边上的计算结果按目的节点的ID进行聚合(A)，输出更新到相应ID的特征。其中S-E， E-E和E-A每一步的输出均可直接用于下一步的计算，而深度学习框架的算子融合并没有覆盖这一融合模式，中间结果的具体化将会导致高内存占用和数据读写耗时。

Seastar设计了一个有限状态机如下图用于识别计算图中可融合的模块并生成融合算子。每个算子的状态由自己的类型标签和上游的算子决定，根据计算图的结构依次遍历每个算子，根据其状态和状态机规则更改其下游算子的状态，遍历结束后，反向记录转换上游合并成融合算子。

为了融合算子的高效执行，需要充分利用硬件的并行机制，将计算与数据合理的分配到并行线程上。对于图神经网络的计算有三个维度可以并行分别为特征并行、节点并行和边并行。其中特征并行是图神经网络区别其他图计算的显著特征，GNN训练过程中计算的最小维度即为节点特征，一般为一个百维至千维的向量。

边并行首先需要将边映射到对应的线程上，当边规模较大时，这一分发过程计算开销将不可忽略，同时对于高入度的节点，在聚合阶段，不同边的计算结果并行写入则会存在冲突。为了提供硬件资源的占用率以及充分利用特征并行，Seastar提出了特征自适应并行策略，此外为了解决实际图数据中存在的高入度节点问题，采用了结合数据本地化执行策略和动态负载均衡的节点并行。

在进行特征并行时，将一个特征向量分配到一个block的线程上，可以使得内存获取连续达成SIMT执行。当特征维度较小，每个block上线程数相应减小则会导致硬件利用率低，block上的线程数大于特征维度则会空载，因而需要进行与特征维度相适应的线程分配。固定block size，增加一个分组维度，一个组可以与其他组共用一个block，也可以占用多个block。

为了实现数据本地化，Seastar为每个节点分配一个分组实现并行，节点的邻边上的计算则是串行，这样在数据本地化上更好，只需载入一次目的节点的特征，在聚合时则将每条边的结果累计更新不需要进行结果同步。不同节点的入度不同，节点并行时则会导致负载不均，为了实现负载均衡，Seastar根据节点入度进行了动态的配置，先根据节点的入度对节点进行排序，这一步可以在数据处理时完成，入度相近的节点计算负载相近被分配到一个block，此外高入度的节点计算时间长优先启动，可以覆盖低入度节点的时间，具体的算子实现模板如下图。

反向计算和前向计算呈现完美对称，前向中心节点接收聚合邻居节点特征，反向中心节点梯度沿着边的反方向发送到邻居节点，因此反向计算可以直接复用前向计算的编译优化，只需把数据的边换个方向，并根据新的方向进行节点排序。为了减少数据的内存占用和获取数据高效，Seastar采用了CSR格式表示图数据，在反向节点排序时需要保持节点ID和边ID与前向一致以便获取正确的前向结果。

对于异构GNN，聚合阶段需要先聚合同类型边，然后聚合不同类型边的结果，具体计算上相比于同构GNN增加了一层外层循环计算边类型，而在融合算子实现上则是将边类型的循环展平，将边类型变动作为一个信号执行外层累加。

结果

以典型GNN网络GAT、GCN、APPNP和R-GCN训练为例来对比评估Seastar与其他框架的性能，相比于DGL-0.4，PyG-1.6.0，单epoch耗时Seastar能实现3-14倍的性能加速，这主要归功于更充分的算子融合和数据本地化更好的执行策略，减少了中间结果的内存读写以及跨block数据获取。在平均入度高的数据集上，Seastar的性能更好，说明了结合数据本地化执行策略和动态负载均衡的节点并行的有效性。在内存占用上，其他框架约为Seastar的2-8倍，主要内存占用在于中间结果为边粒度的特征，相比于节点粒度特征扩大了很多，通过融合优化，Seastar避免了将全量中间结果的具体化。

03

MindSpore Graph Learning

以Seastar的设计为原型，MindSpore Graph Learning前端表达实现了以点为中心的编程模型，进一步改进了通过装饰器来界定以点为中心编程代码，用户可以用For循环遍历中心节点，编程逻辑更接近GNN公式。在代码生成部分，替代tracer和GIR，用parser将以点为中心代码翻译为昇思MindSpore代码，并将翻译结果打印出来，便于用户进行调试。后端执行的算子融合和编译优化部分则是将GNN的融合Pattern下沉到图算融合并复用AKG自动生成高性能算子，其中除了Seastar采用的优化策略，还提出了基于算子融合的重计算方案，解决反向计算依赖前向结果中的超大边粒度特征导致显存不足问题，实现了GNN的高性能训练，相比于业界的框架仍有3-4倍性能提升。

04

总结

近年来图神经网络逐渐发展成为AI一个重要的研究方向，在实际开发应用时，却面临着表达复杂，执行效率低内存占用高等问题。Seastar从前端表达、中间表达和后端执行都提出了一整套的解决方案，在前端表达上以点为中心的编程模型支持用户能用GNN公式逻辑编程忽略计算需要的数据准备，在后端执行上针对GNN的计算模式进行算子融合和优化包括特征自适应的组合并行，数据本地化为中心的节点并行和负载均衡等，实现典型GNN网络训练耗时和内存占用相比其他框架有大幅提升。

对于大规模图数据的分布式训练支持，涉及高效图采样，图切分，不规则通信等难题有待解决，对于异构图神经网络，现有将异构拆分为多个同构执行，忽略了类型本身的信息，精度和性能都有待优化。这些问题都有待图神经网络框架在未来发展演进中解决。