自动化所联合上海交大在图神经网络高效计算方面取得系统性突破

信息时间：2024-07-16 信息来源：中国科学院自动化研究所

近年来，图神经网络（GNN）在许多场景中得到广泛应用，如社交网络分析、推荐系统、自动驾驶等。然而，真实世界中庞大的数据量及图结构显著的不规则稀疏性，为图神经网络的低时延、高能效应用带来了巨大的挑战。随着社会图数据规模的爆炸式增长，解决图神经网络中超大规模数据量带来的访存及计算问题变得越发迫切。

中国科学院自动化研究所程健课题组联合上海交通大学先进体系结构实验室的梁晓峣教授和李钢博士从图神经网络压缩算法设计、专用推理架构设计及高效训练框架设计三方面对图神经网络的高效计算进行了系统和深入地研究，取得系列进展。相关成果相继发表于第11届国际表征学习大会（ICLR 2023）、第30届IEEE国际高性能计算机体系结构大会（HPCA 2024）和第29届ACM架构支持编程语言和操作系统国际会议（ASPLOS 2024），第一作者均为自动化所博士生朱泽雨。



图1. 图数据具有复杂的拓扑结构

在算法设计方面，研发团队提出了端到端的图拓扑自适应的混合量化算法A2Q，在保证模型精度的前提下极大地压缩了图数据。实验结果表明，相比于使用FP32格式存储的图神经网络模型，该压缩方法能够在精度几乎无损的情况下实现高达18.6倍的压缩比。



图2. 研究发现聚合后节点特征的数值范围与其拓扑特性强相关，基于此特性提出了端到端的图拓扑自适应的混合量化算法A2Q

A2Q较高的模型压缩率在降低系统计算延迟和能耗方面展示出了巨大的潜力，然而细粒度的混合精度数据表示和无规则稀疏的节点特征使得该算法在通用处理器上难以充分发挥优势。

为此，在专用推理架构方面，团队针对压缩算法特性和图数据特性设计了一款专用加速器MEGA。通过比特位自适应的存储方式和高效的在线编解码压缩方法，最大限度减少了存储碎片化，降低了数据传输开销。此外，MEGA采用的Condense-Edge调度执行方式，能够将稀疏的边连接聚合在一起以增强访问节点特征时的空间局部性，提升数据复用度。在与英伟达 RTX3090 GPU相当的峰值算力设定下，MEGA实现了121倍的推理加速及476倍的能效提升。



图3. MEGA加速器架构图

更进一步，研发团队提出了全流程图神经网络训练加速系统FastGL，极大地提升了图神经网络在超大规模图数据上的训练速度。与最先进的图神经网络训练加速系统DGL对比，FastGL在多种超大规模图数据和图神经网络模型上实现了平均2.2倍的加速效果，并且具有更好的扩展性，有利于应对未来图数据规模持续增加的趋势。



图4. 全流程图神经网络训练加速系统FastGL训练框架示意图

这一系列图神经网络大规模图数据推理及训练加速方面的研究成果贯穿算法、架构和GPU高效训练框架，形成了较为系统的图神经网络高效计算方案，也为解决通用大模型推理中由超大规模参数量和超长KV cache引起的计算、访存双密集加速等难题提供了启发。

发表该系列研究的国际会议HPCA、ASPLOS与ISCA、MICRO并称为计算机体系结构领域的四大顶会，是计算机体系结构研究的风向标，ICLR则为深度学习领域的顶级学术会议。

论文信息：

[1] Zeyu Zhu, Fanrong Li, Zitao Mo, Qinghao Hu, Gang Li, Zejian Liu, Xiaoyao Liang, Jian Cheng. A2Q: Aggregation-Aware Quantization for Graph Neural Networks. ICLR, 2023

[2] Zeyu Zhu*, Fanrong Li*, Gang Li, Zejian Liu, Zitao Mo, Qinghao Hu, Xiaoyao Liang, Jian Cheng. MEGA: A Memory-Efficient GNN Accelerator Exploiting Degree-Aware Mixed-Precision Quantization. HPCA, 2024

[3] Zeyu Zhu, Peisong Wang, Qinghao Hu, Gang Li, Xiaoyao Liang, Jian Cheng. FastGL: A GPU-Efficient Framework for Accelerating Sampling-Based GNN Training at Large Scale. ASPLOS, 2024

相关开源代码：

A2Q: https://github.com/weihai-98/A-2Q

FastGL: https://github.com/a1bc2def6g/fastgl-ae