[论文笔记] Understanding and Bridging the Gaps in Current GNN Performance Optimizations

最新推荐文章于 2024-10-16 21:42:11 发布

PeakCrosser

最新推荐文章于 2024-10-16 21:42:11 发布

阅读量201

点赞数

分类专栏： Graph Learning 文章标签：论文阅读

本文链接：https://blog.csdn.net/LostUnravel/article/details/133579361

版权

Graph Learning 专栏收录该内容

4 篇文章 0 订阅

订阅专栏

Understanding and Bridging the Gaps in Current GNN Performance Optimizations

理解并弥合当前 GNN 性能优化中的差距 [Paper] [Code] [Presentation]
PPoPP’21

摘要

本文揭示了当前框架在优化 GNN 性能方面的五大差距, 并提出了一套优化措施来填补差距.

1 介绍

GNN 主要由图运算和神经网络运算组成, 但简单地将图计算框架和 DNN 框架组合不足以支持高效的 GNN 执行, 因为其不能有效处理图运算和神经网络运算之间的复杂交互, 而这是 GNN 性能的关键.

本文发现了当前框架在优化 GNN 性能方面的五大差距, 特别是在处理 GNN 相对于传统图或 DNN 操作的特殊复杂性方面. 这些差距存在于数据局部性、负载均衡、冗余(计算)、内存占用和对不同特征长度的处理等方面.

本文提出了一套优化方案来弥补差距. 这些优化利用了最先进的 GPU 优化技术, 并根据 GNN 的特殊性质进行了调整.

局部性感知的任务调度(locality-aware task scheduling)和邻居分组(neighbor grouping): 适应 GNN 中图操作调度, 促进局部性和负载均衡
通过数据可见范围适配器(data visible range adapter)在计算图上调度操作: 解决操作间依赖性, 减少冗余的图数据负载、内核启动、邻居遍历
稀疏提取(Sparse fetching): 避免扩展特征矩阵和冗余计算
调优框架根据给定问题和优化特性选择运行配置

本文贡献:

总结了 GNN 超出传统图计算和 DNN 框架所能处理范围的特殊复杂性.
发现了最先进框架在优化 GNN 性能方面的主要差距, 特别是在处理 GNN 相对于传统图或神经网络运算的特殊复杂性方面.
展示了如何调整 GPU 优化的原理以适应 GNN 的特征, 从而产生了一组显著提高 GNN 计算的优化

2 GNN 优化的复杂性

2.1 GNN 概览

GNN 的两种操作:

图操作: (中心)结点收集邻结点的特征向量, 执行一些操作并更新自身的特征向量
神经网络操作: 在结点间独立执行或根据图结构以中心结点-邻结点的模式执行.

GNN 中的计算层由图操作和神经网络操作组成, 所有计算层构成 GNN 的计算图.

2.2 优化的复杂性

GNN 在性能优化上与传统图计算和 DNN 具有一些共同的复杂性, 但也有一些特性.

2.2.1 共同复杂性

图带来的不规则性
DNN 带来的大量内存占用以及内存访问和计算方面的繁重工作负载.

2.2.2 特殊复杂性

复杂的依赖性和交错模式: 每个 GNN 层中都有各种交织的图运算和神经网络运算.
以图模式执行的神经网络操作: 以中心结点-邻结点模式执行神经网络操作会将图结构的复杂性带到计算中.
不同的特征长度: 带来一些特殊的复杂性, 而高度优化的库无法直接应用, 且图的不规则性和复杂性会被放大.

3 现有 GNN 框架的差距

3.1 方法

3.2 观察到的性能差距

观察 1: 由边粒度带来的图操作局部性较差
一层网络上的图操作可能涉及加载结点特征时 $E * F e a t$ 字节的数据移动 ( $E$ 为图中边数, $F e a t$ 为结点特征的长度), 而结点特征的总大小只有 $N * F e a t$ ( $N$ 为图中结点数), 但现有 GNN 框架很少利用重用 (PyG 边并行).
将中心结点的任务静态分配给底层计算单元, 使内存访问模式完全由图决定, 而图具有不规则的模式 (DGL 结点并行).
在这里插入图片描述

观察 2: 严重的负载不均
现有框架(DGL, ROC, NeuGraph)以中心结点粒度划分任务, 将中心结点及其所有邻居的计算分配给一个计算单元.
图的不规则性使结点的邻居数量可能有很大的差异; GNN 中特征向量长度通常较大, 使得每个相邻结点的计算量很大.
这两个因素导致了严重的负载不均, 从而产生了长尾效应(logn-tail effects).

观察 3: 由密集的函数调用造成的内存访问冗余和开销过大
DGL 和 PyG 使用许多图运算来构建计算图; PyG 方法类似, 将计算划分为许多步骤和内核.
因此, 图结构的重复加载(对于 DGL)或冗余的全局内存访问(对于 PyG)会导致较大的开销.

观察 4: 通过图结构扩展神经网络运算时有大内存占用和冗余计算
以中心结点-邻居结点模式执行的神经网络操作基于图结构, 计算分为两个步骤:

扩展(expansion): 将每个中心结点的邻居特征扩展到连续的内存空间; 会导致较大的内存占用
变换(transformation): 以稠密方式对扩展的矩阵进行神经网络运算, 会导致较大的冗余计算, 即使两个特征向量相同也必须其重新进行变换.

观察 5 :不同特征长度的效率低下
DGL 和 PyG 都不根据特征长度来调整计算.
不同的特征长度使用相同的调度可能会导致性能大幅下降.

4 定制优化以弥补差距

开发了四种优化: 两种针对 GPU 资源和图操作的局部性; 两种针对图操作和神经网络操作的协同优化.

4.1 图操作的任务分配器

局部性感知任务调度(locality-aware task scheduling): 让具有相似邻居分布的中心结点的任务同时运行, 以更好地利用缓存.
邻居分组(neighbor grouping): 每个中心结点的工作负载划分为更细粒度, 来解决负载不均问题.

4.1.1 局部性感知任务调度

相似度度量: Jaccard 相似度; 两个中心结点的 Jaccard 相似度: $\frac{|N_a.neighbor\cap N_b.neighbor|}{|N_a.neighbor\cup N_b.neighbor||}$
调度分为三个步骤:

候选对选择(candidiate pair selection): 找到邻居相似度较高的中心结点对
使用 Min-Hashing 和局部性敏感哈希(Locality-Sensitive Hashing, LSH)将邻结点散列到桶中, 相同桶中的结点更可能相似度高.
对合并(pair merging): 将相似结点合并成更大的簇
将具有高 Jaccard 相似度的结点配对后, 使用优先队列合并相似结点成为更大的簇.
任务调度(task scheduling): 将中心结点调度到实际计算中
将同一簇中结点的任务分配到相邻计算单元, 以更好地重用缓存.

4.1.2 邻居分组

通过设置每个分组中邻居数量的上限(一个可调参数)将每个中心结点的邻居划分为不同分组.
具有大量邻居的中心结点的工作负载将分布到多个计算单元.

带来的问题: 不同计算单元可能映射到不同 SM, 需要进行归约来跨 SM 收集结果.
解决: GNN 运算的大多数归约操作允许在聚合时采取任意顺序, 这允许每个 SM 独立计算归约结果并使用原子指令更新全局内存, 而 SM 间无需数据交换.

4.2 数据可见范围适配器

数据可见范围适配器尝试识别模型中操作的最小可见范围, 同时解决相邻内核可见范围的不匹配问题, 以实现内核融合.
扩展可见范围:

线程到 warp: 利用线程间通信
线程到线程块: 利用共享内存
到全局: 利用运算的线性属性(如 Sum、Mean 和 Div)来推迟所需的全局同步

4.3 稀疏提取和冗余绕过

基本思想: 将神经网络运算与图运算结合起来, 提取通用计算.
稀疏提取(sparse fetching): 根据邻居索引来稀疏访问数据, 将其附加到神经网络操作中, 可将稀疏数据的内存访问从单独的内核移动到神经网络操作内核中.
冗余绕过(redundancy bypassing): 通过限制结点特征矩阵的计算, 利用公共邻居信息来消除冗余变换, 从而将计算量从 $O (E)$ 降低到 $O (N)$ .
在这里插入图片描述