论文标题
减少图形神经网络培训中的沟通
Reducing Communication in Graph Neural Network Training
论文作者
论文摘要
图神经网络（GNN）是使用数据的自然稀疏连接信息的功能强大且灵活的神经网络。 GNN表示这种连接性是稀疏的矩阵，与密集的矩阵相比，算术强度较低，因此沟通成本更高，这使得GNN比卷积或完全连接的神经网络更难扩展到高分子。 我们介绍了用于训练GNN的平行算法系列，并表明与以前的平行GNN训练方法相比，它们可以渐近地减少沟通。我们使用在配备GPU配备的簇上分布的TORCH。实现这些算法，这些算法基于1D，1.5D，2D和3D稀疏密集矩阵乘法。我们的算法优化了整个GNN培训管道的通信。我们在多个数据集上对GNN进行了训练，其中包括一个超过十亿个边缘的蛋白质网络。
Graph Neural Networks (GNNs) are powerful and flexible neural networks that use the naturally sparse connectivity information of the data. GNNs represent this connectivity as sparse matrices, which have lower arithmetic intensity and thus higher communication costs compared to dense matrices, making GNNs harder to scale to high concurrencies than convolutional or fully-connected neural networks. We introduce a family of parallel algorithms for training GNNs and show that they can asymptotically reduce communication compared to previous parallel GNN training methods. We implement these algorithms, which are based on 1D, 1.5D, 2D, and 3D sparse-dense matrix multiplication, using torch.distributed on GPU-equipped clusters. Our algorithms optimize communication across the full GNN training pipeline. We train GNNs on over a hundred GPUs on multiple datasets, including a protein network with over a billion edges.
