arXiv:2411.16127cs.LGcs.PF2024-11被引 4

提升注意力图神经网络在GPU上的训练速度,解决计算碎片化与负载不均问题。

DF-GNN: Dynamic Fusion Framework for Attention Graph Neural Networks on GPUs

  • 动态双层线程调度,灵活适配不同操作,减少冗余计算。
  • 相比DGL稀疏库最高提速7.0倍,端到端训练平均快2.16倍。
  • 兼容PyTorch,适合需要高效训练AT-GNN的科研与工程用户。

注意力图神经网络(AT-GNNs)如GAT和Graph Transformer相较于其他GNN表现更优,但现有图神经网络系统在GPU上训练AT-GNN时效率低下,主要因计算模式复杂。缺乏内核融合导致大量数据移动和显著的内核启动开销;而现有融合策略采用固定线程调度,造成性能瓶颈、冗余计算和负载不均衡。为此,我们提出面向AT-GNN家族的动态内核融合框架DF-GNN。DF-GNN引入动态双层线程调度策略,在保留共享内存优势的同时实现灵活调度,并针对AT-GNN中各类操作定制调度方案,考虑超节点带来的性能瓶颈变化。该框架集成于PyTorch以保证高可编程性。在多种GNN模型与数据集上的评估表明,DF-GNN优于cuGraph和dgNN等现有优化工作,较最先进的非融合DGL稀疏库最高提速7.0倍;端到端训练平均提速2.16倍,显著超越主流DGL框架。

原文摘要 · Abstract (English)

Attention Graph Neural Networks (AT-GNNs), such as GAT and Graph Transformer, have demonstrated superior performance compared to other GNNs. However, existing GNN systems struggle to efficiently train AT-GNNs on GPUs due to their intricate computation patterns. The execution of AT-GNN operations without kernel fusion results in heavy data movement and significant kernel launch overhead, while fixed thread scheduling in existing GNN kernel fusion strategies leads to sub-optimal performance, redundant computation and unbalanced workload. To address these challenges, we propose a dynamic kernel fusion framework, DF-GNN, for the AT-GNN family. DF-GNN introduces a dynamic bi-level thread scheduling strategy, enabling flexible adjustments to thread scheduling while retaining the benefits of shared memory within the fused kernel. DF-GNN tailors specific thread scheduling for operations in AT-GNNs and considers the performance bottleneck shift caused by the presence of super nodes. Additionally, DF-GNN is integrated with the PyTorch framework for high programmability. Evaluations across diverse GNN models and multiple datasets reveal that DF-GNN surpasses existing GNN kernel optimization works like cuGraph and dgNN, with speedups up to $7.0\times$ over the state-of-the-art non-fusion DGL sparse library. Moreover, it achieves an average speedup of $2.16\times$ in end-to-end training compared to the popular GNN computing framework DGL.

图神经网络GPU加速动态调度内核融合

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。