Graph Neural Networks (GNNs) offer a compact and computationally efficient way to learn embeddings and classifications on graph data. GNN models are frequently large, making distributed minibatch training necessary. The primary contribution of this paper is new methods for reducing communication in the sampling step for distributed GNN training. Here, we propose a matrix-based bulk sampling approach that expresses sampling as a sparse matrix multiplication (SpGEMM) and samples multiple minibatches at once. When the input graph topology does not fit on a single device, our method distributes the graph and use communication-avoiding SpGEMM algorithms to scale GNN minibatch sampling, enabling GNN training on much larger graphs than those that can fit into a single device memory. When the input graph topology (but not the embeddings) fits in the memory of one GPU, our approach (1) performs sampling without communication, (2) amortizes the overheads of sampling a minibatch, and (3) can represent multiple sampling algorithms by simply using different matrix constructions. In addition to new methods for sampling, we introduce a pipeline that uses our matrix-based bulk sampling approach to provide end-to-end training results. We provide experimental results on the largest Open Graph Benchmark (OGB) datasets on $128$ GPUs, and show that our pipeline is $2.5\times$ faster than Quiver (a distributed extension to PyTorch-Geometric) on a $3$-layer GraphSAGE network. On datasets outside of OGB, we show a $8.46\times$ speedup on $128$ GPUs in per-epoch time. Finally, we show scaling when the graph is distributed across GPUs and scaling for both node-wise and layer-wise sampling algorithms.

翻译：图神经网络（GNN）提供了一种紧凑且计算高效的方式来学习图数据的嵌入表示与分类。GNN模型通常规模庞大，因此需要进行分布式小批量训练。本文的主要贡献在于提出了降低分布式GNN训练中采样步骤通信开销的新方法。我们提出了一种基于矩阵的批量采样方法，将采样过程表示为稀疏矩阵乘法（SpGEMM），并一次性采样多个小批量。当输入图拓扑结构无法容纳于单个设备时，我们的方法通过分布式存储图结构并采用通信避免型SpGEMM算法来扩展GNN小批量采样，从而支持在远超单设备内存容量的超大规模图上进行GNN训练。当输入图拓扑结构（不含嵌入表示）可容纳于单块GPU内存时，本方法实现了：（1）无通信采样；（2）均摊小批量采样的额外开销；（3）仅通过不同的矩阵构造即可表征多种采样算法。除新提出的采样方法外，我们设计了一套基于矩阵批量采样的端到端训练流水线。在包含128块GPU的最大规模开放图基准（OGB）数据集上的实验表明，对于3层GraphSAGE网络，我们的流水线比Quiver（PyTorch-Geometric的分布式扩展）提速2.5倍。在OGB之外的测试数据集中，单轮训练时间在128块GPU上实现了8.46倍的加速。最后，我们展示了方法在跨GPU分布式图场景下的扩展能力，并验证了其同时适用于节点级采样与层级采样算法。