面向图神经网络训练的分布式基于矩阵的采样

图神经网络（GNN）提供了一种紧凑且计算高效的方式来学习图数据上的嵌入与分类。GNN 模型通常较大，使得分布式小批量训练成为必要。本文的主要贡献是用于减少分布式 GNN 训练采样步骤中通信的新方法。在此，我们提出一种基于矩阵的批量采样方法，将采样表示为稀疏矩阵乘法（SpGEMM）并一次性采样多个小批量。当输入图拓扑无法放入单个设备时，我们的方法对图进行分布并使用避免通信的 SpGEMM 算法来扩展 GNN 小批量采样，使得 GNN 训练能够在比单个设备内存所能容纳的图大得多的图上进行。当输入图拓扑（而非嵌入）可放入一个 GPU 的内存时，我们的方法（1）在无通信的情况下执行采样，（2）摊还采样一个小批量的开销，以及（3）可仅通过使用不同的矩阵构造来表示多种采样算法。除新的采样方法外，我们引入了一个使用我们基于矩阵的批量采样方法以提供端到端训练结果的流水线。我们在最大的 Open Graph Benchmark（OGB）数据集上基于 $128$ 个 GPU 给出了实验结果，并表明我们的流水线在 $3$ 层 GraphSAGE 网络上比 Quiver（PyTorch-Geometric 的分布式扩展）快 $2.5\times$。在 OGB 之外的数据集上，我们在 $128$ 个 GPU 上展示了每轮时间 $8.46\times$ 的加速。最后，我们展示了当图分布在 GPU 之间时以及针对逐节点和逐层采样算法的扩展性。