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 show that judiciously replicating feature data with a simple all-to-all exchange can outperform current methods for the feature extraction step in distributed GNN training. 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 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训练中减少采样步骤通信开销的新方法。我们提出了一种基于矩阵的批量采样方法，将采样过程表达为稀疏矩阵乘法（SpGEMM），并一次性完成多个小批量的采样。当输入图拓扑无法容纳于单个设备时，该方法通过分布式存储图结构并采用通信避免型SpGEMM算法来扩展GNN小批量采样规模，从而支持在远大于单个设备内存容量的图上进行GNN训练。当输入图拓扑（而非嵌入表示）可适配于单个GPU内存时，本文方法能够：（1）无需通信即可完成采样；（2）分摊小批量采样的开销；（3）通过简单地采用不同矩阵构建方式实现多种采样算法。除新型采样方法外，我们还证明通过简单的全交换策略对特征数据进行合理复制，可以在分布式GNN训练的特征提取步骤中优于现有方法。我们在128个GPU上对最大规模开放图基准（OGB）数据集进行了实验验证，结果表明所提流水线在3层GraphSAGE网络上比Quiver（PyTorch-Geometric的分布式扩展）快2.5倍。在OGB之外的测试数据集上，我们实现了每个训练周期在128个GPU上达到8.46倍的加速比。最后，我们展示了当图结构跨GPU分布式存储时的扩展能力，以及节点采样与层级采样两种算法的扩展性能。