Communication Optimization for Distributed Execution of Graph Neural Networks
Communication Optimization for Distributed Execution of Graph Neural Networks
复制标题
DOI：
10.1109/ipdps54959.2023.00058
复制
发表时间：
2023-05
期刊：
影响因子：
--
通讯作者：
Süreyya Emre Kurt;Jinghua Yan;Aravind Sukumaran-Rajam;Prashant Pandey;P. Sadayappan
中科院分区：
文献类型：
--
作者：
Süreyya Emre Kurt;Jinghua Yan;Aravind Sukumaran-Rajam;Prashant Pandey;P. Sadayappan
图形神经网络(GNN)已经成为一种非常强大和流行的机器学习模型，适用于许多应用领域。GNN的每个阶段都需要聚集(稀疏矩阵-矩阵乘法)和线性运算(密集矩阵-矩阵乘法)。许多努力都涉及GNN分布式实施的发展问题。虽然分布式矩阵乘法的高效算法已为人所熟知，但GNN所需的分布式矩阵乘法序列的集体优化是一个挑战，其中分量矩阵乘法运算的排序也存在许多自由度，本文提出了一种新的分布式GNN方法--矩阵重分布(RDM)，其核心是通过GNN级间的矩阵重分布实现无通信的分布式矩阵乘法。虽然该方法适用于GNN的众多算法变体，但实验评估集中在GCN(图卷积网络)上，包括全批训练和使用GraphSAINT的基于采样的训练。在一个具有8个GPU的多GPU系统上，使用128或256个隐藏特征，对2层和3层GCN进行了实验评估，结果表明，RDM在CAGNET和DGCL两种最先进的多GPU GCN实现中获得了2×到3.7倍的几何平均加速比。
Graph Neural Networks (GNNs) have emerged as a very powerful and popular machine learning model for numerous application domains. Each stage of a GNN requires an aggregation (sparse matrix-matrix multiplication) and a linear operation (dense matrix-matrix multiplication). Numerous efforts have addressed the development of distributed implementations for GNNs. Although efficient algorithms for distributed matrix multiplication are well known, the challenge here is the collective optimization of sequences of distributed matrix-matrix multiplications required for GNN, where many degrees of freedom also exist in the ordering of the component matrix-multiplication operations.This paper develops a new approach to distributed GNN, ReDistribution of Matrices (RDM), centered around communication-free distributed matrix-multiplication enabled by matrix redistribution between GNN stages. While the approach is applicable to the numerous algorithmic variants of GNN, the experimental evaluation focuses on GCN (Graph Convolutional Network), including both full-batch training as well as sampling-based training using GraphSAINT. Experimental evaluation with 2-layer and 3-layer GCN, using 128 or 256 hidden features, across eight sparse datasets, on a multi-GPU system with 8 GPUs shows that RDM attains a geometric mean speedup between 2× and 3.7× over two state-of-the-art multi-GPU GCN implementations, CAGNET and DGCL.