GLPilot：利用可学习嵌入实现的高效分布式图神经网络（GNN）训练

《IEEE Transactions on Parallel and Distributed Systems》：GLPilot: Efficient Distributed GNN Training With Learnable Embeddings

　　嵌入数据并行训练优化方法及系统实现

　　

摘要：

图神经网络（GNNs）通过可学习的顶点嵌入，使得模型即使在顶点特征稀疏、有噪声或缺失的情况下，也能推断出丰富且与任务相关的表示。在大规模的多GPU训练中，动态更新的嵌入（其大小通常比模型参数大几个数量级）会严重降低训练效率。具体来说，加载远程嵌入以及同步它们的梯度占每次迭代时间的90%以上。传统的缓存和并行处理方法是为静态嵌入或模型参数单独设计的，对于缓解这种与嵌入相关的传输“数据墙”问题效果不佳。为了解决这个问题，我们首先详细分析了训练过程中的顶点访问模式，发现虽然频繁被访问的顶点占据了大部分嵌入加载延迟，但它们的更新次数却很少，因此这些顶点的嵌入非常适合被重复使用（即“陈旧数据”的再利用）。基于这一发现，我们提出了GLPilot这一创新系统，用于缓解与嵌入相关的瓶颈。GLPilot引入了一种限制“陈旧性”的嵌入缓冲机制来减少远程数据获取次数，并采用了一种局部梯度聚合技术来最小化同步过程中的冗余通信。此外，GLPilot还利用GPU上的缓存来存储最新更新的嵌入数据，从而缓解CPU与GPU之间的数据传输瓶颈。我们在一个32核GPU集群上使用两种流行的GNN模型、三个数据集和两种优化器进行了测试，结果表明，与DGL和P3这两个强大的基准算法相比，GLPilot能够将每次迭代的训练速度提高1.28至1.93倍，同时保持模型精度相当。