arXiv:2504.04673cs.LG2025-04被引 8

针对图神经网络训练通信瓶颈,提出感知稀疏性的通信优化方法。

Sparsity-Aware Communication for Distributed Graph Neural Network Training

  • 仅传输必要矩阵元素,避免冗余通信
  • 通过图划分重排矩阵,通信量降低至接近零
  • 适配多卡训练,适合大规模GNN分布式场景

图神经网络(GNN)在图数据上学习嵌入和分类具有计算高效性。然而,其训练计算强度低,通信开销成为可扩展性的瓶颈。全图训练中的核心操作是稀疏矩阵与稠密矩阵乘法(SpMM)。以往并行化方法采用忽略稀疏性的算法,无论稀疏模式如何均传输所有矩阵元素,虽有可预测的通信模式且支持集体通信,但浪费大量带宽。本文提出三种感知稀疏性的新方法:一是仅通信必需元素;二是利用图划分模型重排矩阵,大幅减少传输量;三是通过定制化划分缓解通信负载不均衡,同时最小化总通信量与最大发送量。进一步将这些方法与通信规避策略(1.5D并行SpMM)结合,通过子矩阵复制降低通信。实验表明,在256张GPU上最高提升14倍,部分实例实现近乎零通信的并行训练,相较主流基于通信无关SpMM的框架显著优化。

原文摘要 · Abstract (English)

Graph Neural Networks (GNNs) are a computationally efficient method to learn embeddings and classifications on graph data. However, GNN training has low computational intensity, making communication costs the bottleneck for scalability. Sparse-matrix dense-matrix multiplication (SpMM) is the core computational operation in full-graph training of GNNs. Previous work parallelizing this operation focused on sparsity-oblivious algorithms, where matrix elements are communicated regardless of the sparsity pattern. This leads to a predictable communication pattern that can be overlapped with computation and enables the use of collective communication operations at the expense of wasting significant bandwidth by communicating unnecessary data. We develop sparsity-aware algorithms that tackle the communication bottlenecks in GNN training with three novel approaches. First, we communicate only the necessary matrix elements. Second, we utilize a graph partitioning model to reorder the matrix and drastically reduce the amount of communicated elements. Finally, we address the high load imbalance in communication with a tailored partitioning model, which minimizes both the total communication volume and the maximum sending volume. We further couple these sparsity-exploiting approaches with a communication-avoiding approach (1.5D parallel SpMM) in which submatrices are replicated to reduce communication. We explore the tradeoffs of these combined optimizations and show up to 14X improvement on 256 GPUs and on some instances reducing communication to almost zero resulting in a communication-free parallel training relative to a popular GNN framework based on communication-oblivious SpMM.

图神经网络通信优化稀疏性分布式训练

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