最近，图神经网络(GNNs)已经成为分析非欧几里得图数据的最新算法。然而，实现高效的GNN训练是一项挑战，特别是在大规模图上。本文基于一篇于2022年发表在PACT上的论文，主要介绍基于近内存NMP的大规模图神经网络训练加速器GNNear。
首先介绍GNN和GNN训练的基本知识。一般来说，GNN包含多层GNN网络。如图所示，给定的输入图G包含图的顶点（如A,B,C等）和边（如A和C的连接）。GNN通过如上图所示的两个步骤进行计算
上式中，l表示GNN的层数。在这一过程中，首先使用Aggregate步骤收集顶点v的邻接点特征hul-1。接着，Combine步骤利用神经网络处理Aggregate的结果avl 。
为了从图中学习有用的信息，GNN必须经历训练过程。其中Combine步骤的后向传播与神经网络相同，Aggregate步骤的后向传播需要沿着边聚集每个顶点的特征梯度。具体公式如下所示。
因此，GNN的训练过程可抽象为两个不同的计算，即Reduce与Update。
接着介绍GNN大规模训练的特点。
(1) GNN训练会产生很大的内存占用。如图1所示，GNN应用图的规模不断增大。同时，反向传播需要大量的缓存数据。例如，使用PyG框架对包含1.14亿条边的中等规模的Reddit数据集进行训练，需要58GB的内存保存中间数据，并产生超过300GB的DRAM访问。
(2) GNN训练既有内存密集型又有计算密集型。Reduce操作涉及特征的聚集，属于内存密集型。Update操作涉及全连接层神经网络的计算，属于计算密集型。如图4所示，Reduce操作与Update操作是GNN训练过程中的主要瓶颈。
(3) GNN需要考虑图的不规则结构，以此优化资源利用率和解决负载不平衡问题。
为了克服上述挑战，本文提出了一个混合的GNN训练加速器GNNear，结合了基于DIMM的近内存处理引擎(NME)和一个集中式加速引擎(CAE)。
集中加速引擎主要负责计算密集型的Update操作，并还合并由近内存处理引擎产生的部分和。集中式加速引擎包含一个GEMM引擎（用收缩阵列实现）和一个矢量处理单元(VPU)。该控制器可根据输入图的相邻矩阵和模型配置安排整个训练过程。近内存处理引擎位于DDR4的缓冲芯片中。
每个支持NMP的DIMM都包含一个NME。NME主要包含指令寄存器、指令解码器、执行单元、数据缓冲器和控制器。NME从CAE接受NMP指令并放入指令寄存器。指令解码器对指令进行解码。接着，轻量级的控制器根据指令进行执行。执行单元处理数据计算。其中，控制器还可生成常规的DDR4 Command/Address和data信号。
上图描述了架构的基本工作流程。
1. 首先，划分顶点并将他们平均分配给不同的DIMM存储（DIMM0与DIMM1），如图6-(a)。
2. 对于Reduce: CAE通过内存接口向NME发送定制指令，NME在本地进行计算，并由CAE读出（操作1）。如图6-(b)左侧所示，描述了图的邻接矩阵，其中的source vertexes是图G的顶点，表示当前GNN层的输入；destination vertexes表示当前GNN层的输出。以图中计算的节点5为例，Y5 = a0′+a1′，其中a0′ = X1+X4，a1′= X5+X6。NME将部分和a0′与a1′传回CAE，CAE将计算结果合并(Y5 = a0′+a1′)。
3. 对于update：CAE从片上缓存或直接读取数据，必要时将数据写回DRAM（操作2）。
4. 因为操作之间没有依赖，所以可以用流水线实现，如图6-(c)。
为了进一步减少DRAM的访问，利用图形分片策略探索数据重用。如图8所示，DIMM0的计算了目标顶点V1~V8的部分和。我们把邻接矩阵分为多个R×C大小的切片每一步处理一个切片。R×C = 1×2中step9的计算为例，首先加载源顶点X4（操作1）。然后，在同一 片内进行计算。其中，块2和块3的计算都需要X4。产生的部分和Y5’与Y6’也可被后续的切片重用。
本文设计了三类指令，用于近内存加速工作使用，如图10所示。其中，L-Type从DRAM中加载顶点数据；R-Type从NME中取出计算数据（不通过DRAM）；C-Type控制近内存计算。Op确定Aggregate步骤的方式，Edge_w是权重数据，Dst_存储部分和数据。
为了处理图的不规则性，首先提出了HGP分区方法。分区策略需要将源顶点的值分配到不同的DIMM中，每个DIMM通过访问存储的源顶点计算目标顶点V1~V8的特征值。如图13所示，Channel0包含了4个DIMM，4个DIMM以不同的颜色深浅表示。考虑到高度顶点（有较多相连的顶点）与低度顶点的特征，即高度顶点的内存访问需求高于低度顶点。DIMM中高度顶点的存在会带来大的负载。并且，在数据集中，低度顶点的占比往往大于高度顶点。低度顶点的邻接点少，需要在NME中Reduce的需求少，会带来资源利用不足的情况。
为此，本文提出的HGP划分，将低度顶点均匀分配在不同DIMM上，计算所有目标顶点的值。将高度顶点复制到所有DIMM中，并计算部分目标顶点的值。可以看到，经过HGP，在每个DIMM中的工作负载是基本一致的。
接着介绍架构的调度优化。本文提出基于窗口的调度，以解决图的不规则带来的问题。我们将不同的目标顶点的计算分割为不同的间隔，例如，在间隔1中处理目标顶点V1~V8。
1. 如果间隔𝑖+1在处理窗口内(in-Processing)，我们允许CAE端控制器在内存完成间隔𝑖后立即对间隔𝑖+1发出指令。
2. fifo中的部分结果将与存储在窗口缓冲区(window buffer)中的部分结果合并。
3. 一旦每个DIMM的间隔𝑖的结果被合并，如图15中的间隔5，CAE就会提交间隔𝑖并右移处理窗口。通过这种方法，我们可以同时调度多个区间，并减轻由图的不规则性引起的负载不平衡问题。
本文使用定制的Dramsim3模拟器进行评估。并在多个数据集与多种GNN算法与Xeon E5-2698-V4CPU,V100 GPU比较吞吐量与能效的提升。如图所示，GNNear实现了30.8/2.5的速度提升和79.6/7.3的性能提升。
 关 注 我 们 
实验室网站：http://viplab.fudan.edu.cn/
OpenASIC官方网站：www.openasic.org
知乎专栏：http://zhuanlan.zhihu.com/viplab
微信公众号：OpenASIC
继续阅读
阅读原文