图模型训练论文笔记

  |   0 评论   |   0 浏览

1. Deep Graph Library: A Graph-Centric, Highly-Performant Package For Graph Neural Networks

【摘要】

提出了DGL的设计原则和具体实现,提取了GNN的模型计算范式,将其抽象成了通用的矩阵计算,目的是更好地实现并行化。优势是对比现有框架减少了内存占用同时提高了训练速度。

【Introduction】

总结了GNN的使用前景并提出需要一个 对研究人员而言灵活有理并且对实际应用而言高效能用的工具。

两个主要挑战:

  1. 现有的以张量为核心的深度学习框架与图之间存在语义鸿沟。
  2. 图的计算/存储的稀疏性和并行硬件,dense OP之间存在性能差异。

主要贡献:

  1. 将GNN的计算步骤提取成一个用户可配置的消息传递模型。概括了前向和反向的稀疏矩阵运算用来进行硬件加速,同时探索了在此基础上的一些加速策略。
  2. 制定了以图为中心的编程抽象,主要表现在其代码库内将整个图的操作封装成了DGLGraph和DGLHeteroGraph类中,对外只暴露接口,底层可根据不同硬件进行不同的实现。
  3. 将GNN除模型以外的部分解耦,向下可以兼容Tensorflow / pytorch, Mxnet等多个框架。

【图神经网络和消息传递】

介绍了一下消息传递范式:将整个图的更新分为点上更新和边上更新两个步骤,对应消息函数,聚合函数,更新函数三个函数。

假设图可以表示成 \mathcal{G}(\mathcal{V}, \mathcal{E}), x_v \in R^{d_1} 表示节点 v 的特征,w_e \in R^{d_e} 表示边 (u, v, e) 的特征。

msgpass.png

消息传递包含边上传递和节点传递两个阶段,边上传递的过程由消息函数 \phi 根据边和点的特征计算 message。每个节点有一个mailbox可以接收边传来的message。之后聚合函数 \rho 聚合 message 并由更新函数 \psi 更新节点的特征。

GNN中,通常 \phi\psi 为神经网络,\rho 为 sum/mean/max(min)/LSTM 等pooling函数。

【通用的SpMM和SDDMM】

消息传递和稀疏矩阵操作有很强的联系。本节总结了图模型计算中的两种矩阵运算的抽象,针对这两种不同的抽象可以进行硬件加速的优化以及接口的定义。

1. SPMM

指计算节点特征时稀疏矩阵和稠密矩阵相乘的操作(Sparse dense matrix multiply)。例如给定一个节点的特征矩阵 X \in R^{|v| \times d} 和图的邻接矩阵 A ,节点级的GCN 可以表示成一个sparse矩阵和一个dense矩阵相乘Y=AX

generalized SpMM表示如下:

gSpMM_{\mathcal{G},\phi(z),\rho}: \mathbb R^{|\mathcal V|\times d_1},\mathbb R^{|\mathcal V|\times d_2},\mathbb R^{|\mathcal E|\times d_3} \mapsto \mathbb R^{|\mathcal V|\times d_4}
Z=gSpMM_{\mathcal{G},\phi(z),\rho}(X, Y, W)

SpMM由节点和边的特征输出节点的表示。

2. SDDMM

指计算边特征时的矩阵运算规则。例如很多GNN模型会计算边的权重。通常可以抽象为一个权重矩阵与邻接矩阵的点积,通常权重可由特征矩阵的点乘表示。

W=A\odot(XX^T)

SDDMM可以看成首先是两个Dense矩阵相乘,之后再与Sparse矩阵求点积。

generalized SDDMM表示如下:

gSDDMM_{\mathcal{G},\phi(z),\rho}: \mathbb R^{|\mathcal V|\times d_1},\mathbb R^{|\mathcal V|\times d_2},\mathbb R^{|\mathcal E|\times d_3} \mapsto \mathbb R^{|\mathcal E|\times d_4}
M=gSDDMM_{\mathcal{G},\phi(z),\rho}(X, Y, W)

SpMM由节点和边的特征输出节点的表示。

3. 与消息传递、并行的关系

SPMM对应节点级的更新,也就是聚合函数和更新函数。SDDMM对应边上信息的更新,也就是message的计算。通俗的理解就是所有GNN的操作都可以抽象成矩阵之间的运算,按照输出数据的维度可以分为不同的两种抽象。做这两种抽象的原因是对于这两种计算选取的图的存储方式以及放到GPU上运行的方式可以不同,下表显示了这种不同。DGL现在采用了一种启发式的策略去执行这两种并行,通常g-SpMM使用节点并行的方式,g-SDDMM采用边并行的方式。

dglparallel.png

【DGL系统设计】

1. Graph as a first-class citizen

  1. 以Graph为中心,使用user-friendly的抽象,允许底层深度优化。
  2. 框架中立性,通过接Adapter的方式可以实现无缝的应用迁移。
  3. 向下屏蔽了图存储的一些细节:CSR/CSC分别用于前向和反向的g-SpMM, COO用于g-SDDMM。
  4. 用户通过g.update\_all(\psi, \rho)g.apply\_edges(\phi)去触发g-SpMM和g-SDDMM的kernel,编译器会解析出给定的函数实现一个fused_kernel。
  5. message reduce阶段会按节点的degree分bucket,每个bucket中的msg可以看成一个dense tensor。

2. Framework-neutral design(框架中立性)

  1. 保留了对图模型计算中最关键的部分控制与实现,例如稀疏矩阵的管理与操作。
  2. 使用其他框架的tensor作为输入,并使用他们的dense OP及操作。DGL定义了一些shim(楔子)将其使用的 dense OP 映射到指定框架的实现,类似ONNX。

【实验】

GPU 机器: 1 NVidia V100 GPU with 16GB GPU RAM and 8 VCPUs
CPU 机器: 64 VCPUs and 256GB RAM

对比了不同数据集上DGL和PyG、GraphNets的性能差异,得到DGL在训练速度和内存消耗上都有较大优势。

dglexp.png

2. Efficient Data Loader For Fast SAMPLING-BASED GNN Training On Large Graphs

【摘要】

提出了单机多卡上基于采样设计的DataLoader,核心是利用可用的GPU资源做Data的Cache,使用了一种轻量且有效的Cache Policy。
提出了一种GNN-computational-aware的图partition算法,算法提高了cache的有效性,避免了cross-partition访问。

【问题描述】

通过实验的方式比较了Neighbor-Sampling(每个节点取固定个数的邻居)和Layer-wise Sapling(每一层Layer取固定数目的邻居)两种采样方式下训练数据加载时长占总时长的比例。

总结了现有DGL框架的几个问题:

  1. 数据加载时间占整个训练时间的绝大部分。
  2. 节点冗余访问问题。(存在一些热节点)。
  3. GPU利用率不高(数据加载时间和训练时间不匹配导致的)。
  4. 数据加载和GNN的计算是线性的过程,没有流水线的操作。

【PA-GRAPH】

总体架构如下图所示:

pagraphach.png

提出了PA-Graph框架及其3个核心技术:

  1. GNN亲和的Cache机制,减少 CPU->GPU 的数据拷贝。
  2. 提出了一个Cache友好的数据并行训练方法。
  3. 数据加载和训练的overlap。

1. GNN Computation-aware caching

(1) Cache policy

如果使用常规的对训练样本random-shuffle的方式,不能预测每个mini-batch中的节点,因为邻居是随机选择的,所以很难预测哪些节点会在下一个mini-batch被拿到。

所以提出了使用节点的out-degree作为节点被select可能性的标准,理由是一个节点的out-degree越大,该节点更有可能作为其他节点的in-neighbor。文章同时认为动态Cache的方式不适合GNN的训练,因为cache数据是多个线程共享的,所以PA-GRAPH使用的cache策略非常简单,就是按照节点out-degree进行排序的静态cache。

(2) Cache 大小

文中指出为了更好地利用宝贵的缓存,在不影响训练的前提下Cache的大小应该尽可能的大。所以在第一个mini-batch结束后可以得到可用的Cache大小。

2. 训练和Partition之间的数据并行。

文中指出DGL在多卡训练时采用的是模型并行的方式,也就是多个GPU worker拿到的是同一份图数据,这样做cache利用率很低,不能充分利用图访问的空间局部性,所以PA-GRAPH采用了数据并行的方法,每个卡负责自己的partition,多个卡之间交换相互的梯度。以此为目标设计了一个图的partition算法。算法设计的目标主要考虑两方面(1)Trainer之间负载均衡。(2)减少跨partition之间的重叠,在分partition的时候会同时存储边界节点的邻居,这些邻居节点不参与训练,目的是不去查全局的store_server,如果这种重叠多了会造成存储浪费。

提出了一种贪心算法:总得partition数量固定的情况下,每个节点加入partition之前会计算一个score,加入score最高的partition中。

Score = |TV_i \cap IN(V_t)|\cdot \frac{TV_{avg}-|TV_i|}{|PV_i|}

TV_i表示partition i中当前训练节点的数目, IN(V_t) 表示该节点n-hop的邻居有几个在当前partition中,TV_{avg}表示每个partition应有节点数的期望值,|PV_i| 表示partition i中的总结点数,包括训练节点数和其n-hop内的邻居。

3. 数据加载和GNN训练的pipeline

整个流分为Loading和Computation,由MsgQueue进行交互和同步。

pagraphpipeline.png

3. BGL:GPU-Efficient GNN Training by Optimization Graph Data I/O and Preprocessing

【摘要】

指出现有图模型训练系统不够高效的原因是为GPU准备数据态慢,该部分过程包括了子图的采样和特征提取。

提出了BGL框架,为了解决现有瓶颈,主要工作包括:
(1)设计了一个动态Cache引擎,减少特征检索的开销。
(2)改进了图的partition算法,减少了子图采样的跨partition之间的通信。
(3)通过资源隔离减少了数据预处理阶段的资源竞争。

【Introduction】

场景:图规模20亿节点,2万亿边。按图采样每次取mini batch进行训练。每次迭代可以分为如下3个阶段:
(1)在分布式图存储服务器采样子图。
(2)子图的特征检索。
(3)GNN的前向和反向计算。

指出了现有框架(DGL,PyG,Euler)存在的一些问题:
(1)取得训练样本需要大量数据传输。
(2)GPU利用率不高,DGL,PyG虽然采用了数据预取的方式,但在large Graph上只有10%左右的GPU利用率。

BGL通过动态cache机制,图分隔算法改进和优化资源隔离的方式相较于PaGraph, PyG, DGL 和Euler分别有2.14x,3.02x,7.04x和20.68x的加速比。在使用远程分布式存储的情况下依然可以使V100的GPU打满。

【设计】

Feature Cache Engine

1. Cache算法

指出PaGraph中提出的静态Cache机制不适合大规模的图场景,所以Cache算法采用动态Cache的方式。首先调研了几种动态Cache的方式,FIFO的性能 > LRU 或 LFU。提出采样顺序具有空间局部性,所以可以采用多源BFS的顺序进行采样。但该方法同样存在trade-off,会造成模型收敛变慢。

所以在bfs的基础上加上了随机性的因素,使用多个BFS队列,随机选取若干个队列,并在采样过程中随机置换出其中一些队列。下图显示了这种双随机采样方式的流程,每个batch会从多个队列中取节点作为training node,但对其中的负采样方法并未体积,估计是取与当前节点队列不同的节点作为负节点。

bglsample.png

2. Cache设计

Cache设计采用CacheMap+CacheBuffer的方式(CPU+GPU双级Cache,均采取这种方式),其中CacheMap中存指针,CacheBuffer中取具体的BufferSlot,分布式存储采用分桶的方式,每个GPU中的CacheMap管理自己的CacheBuffer,多卡之间的CacheBuffer通过NvLink连接。

3. Cache的workflow和一致性保证

整个Cache的步骤如下

(1)将子图节点按ID进行分桶,将查询任务加入到一个CacheQueryQueue(为了防止对GPU的访问出现线程竞争的现象,BGL将这种访问都放到一个访问队列,有一个单独的线程专门负责,相较于使用锁减少了8倍的overhead)。
(2)对CacheQueryQueue分配线程。
(3)通过CacheMap中查到指针进而访问CacheBuffer。
(4)如果GPU Cache Miss会去查CPU Cache,如果再Miss会去查RemoteServer。
(5)如果查的是RemoteServer,结果会立即送到GPU,随后根据Cache替换规则更新CacheBuffer和CacheMap。

Graph Partition

首先指出在大规模图训练场景下一个好的Partition算法应该具有如下特征:
(1)可以拓展到10亿节点的规模。
(2)一个节点多跳节点可以在一个或者很少量的请求中获得。
(3)各个Partition之间training Node应该负载均衡。

本文提出的Partition算法两大创新点:
(1)多级粗粒度的策略减少partition算法的复杂度,同时保持多跳之间的连接性。
(2)提出了一种启发式的新颖的节点分配方式,在保证连接性的同时保证训练节点的负载均衡。

bglpartition.png

1. partition 模块 work-flow

整个Partition的workflow包含如下一些Step:

  • Step1:Coarsening
    多源BFS将图切分成多很的blocks,使用一种random-partition算法从HDFS加载数据。然后随机选一些源节点分配的blockId,这些源节点可以通过BFS广播自己的blockId。
  • Step2: Assignment
    利用下文提出的启发式partition算法将blockId分配到不同的partition上。
  • Step3:Uncoarsening
    利用blockId -> PartitionId的映射,将block映射到nodes,然后将映射好的nodes送往HDFS。

2. 启发式分配算法

采用一种贪心策略进行节点partition分配,综合考虑:
(1)加入partition的局部连接性。
(2)partition内节点数目。
(3)跨partition交互的权重。

bglpartitionalgo.png

k 是 partition 号,P(i) 表示第 i 个partition,\Gamma^j(B)表示该block j hop的所有邻居,T(i)表示第 i 个partition中训练节点的个数,CC_T 均为惩罚系数。感觉跟PaGraph差不多。

Resource Isolation For Contending Stages

主要解决训练过程中对CPU的争用。将GNN训练过程分为了一个8级的异步流水线。

bglpipeline.png

在流水线基础上考虑了(1)StoreServer上处理采样请求和生成子图的CPU竞争。(2)子图处理和Cache机制对WorkerServer的CPU竞争。(3)子图和Feature向GPU拷贝对PCIE的竞争。

提出了一个最优化问题,解决之后可以算出 (1),(2) 中各自占用的CPU核心数和 (3) 中所占的PCIE带宽。

【实现】

重用了DGL的图存储模块,利用GMiner的图处理模块进行Partition,和DGL团队合作进行上游开发。

FeatureCacheEngine

(1)使用一个单独的CudaStream去执行Cache逻辑。
(2)CPU使用OpenMP加速FIFO的操作。

IPC

使用共享内存避免IPC OverHead,Linux端使用了Linux SharedMem, Cuda使用了Cuda的IPC库。


标题:图模型训练论文笔记
作者:YaoCheng8667
地址:https://ycisme.xyz/articles/2022/03/25/1648201464382.html