面向GPU的注意力图神经网络动态融合框架DF-GNN
机器学习
2024-11-26 v1 性能
摘要
注意力图神经网络( AT-GNN)，如GAT和Graph Transformer，已显示出优于其他图神经网络的卓越性能。然而，现有图神经网络系统在GPU上高效训练AT-GNN面临计算模式复杂的挑战。在未进行内核融合的情况下执行AT-GNN操作会导致重大的数据移动和显著的内核启动开销，而现有图神经网络内核融合策略中固定的线程调度会导致次优性能、冗余计算和不均衡的工作负载。为此，我们提出一种名为DF-GNN的动态内核融合框架，用于AT-GNN系列。DF-GNN引入动态双层线程调度策略，使线程调度能够灵活调整，同时保留融合内核中共享内存的优势。DF-GNN为AT-GNN中各操作量身定制特定的线程调度，并考虑到超节点的存在导致的性能瓶颈迁移。此外，DF-GNN集成于PyTorch框架以实现高可编程性。跨越多种图神经网络模型和多个数据集的评估表明，DF-GNN超过了诸如cuGraph和dgNN等现有图神经网络内核优化工作，在对最新非融合DGL稀疏库上实现最高可达7.0倍的加速。更重要的是，它在流行图神经网络计算框架DGL上的端到端训练平均加速达2.16倍。
引用
@article{arxiv.2411.16127,
title = {DF-GNN: Dynamic Fusion Framework for Attention Graph Neural Networks on GPUs},
author = {Jiahui Liu and Zhenkun Cai and Zhiyong Chen and Minjie Wang},
journal= {arXiv preprint arXiv:2411.16127},
year = {2024}
}