[论文解读] DistGNN: Scalable Distributed Training for Large-Scale Graph Neural Networks.

Vasimuddin, Sanchit Misra+7 more|arXiv (Cornell University)|Apr 14, 2021

Advanced Graph Neural Networks参考文献 34被引用 13

一句话总结

DistGNN 通过优化深度图库（DGL）以提高共享内存效率、采用最小顶点割图划分以减少通信，并引入延迟更新算法以避免通信，提出了一种可扩展的分布式训练框架，用于在 CPU 集群上进行大规模图神经网络（GNN）训练。与单插槽 CPU 训练的基线 DGL 相比，该框架在使用 128 个 CPU 插槽时实现了高达 97 倍的加速。

ABSTRACT

Full-batch training on Graph Neural Networks (GNN) to learn the structure of large graphs is a critical problem that needs to scale to hundreds of compute nodes to be feasible. It is challenging due to large memory capacity and bandwidth requirements on a single compute node and high communication volumes across multiple nodes. In this paper, we present DistGNN that optimizes the well-known Deep Graph Library (DGL) for full-batch training on CPU clusters via an efficient shared memory implementation, communication reduction using a minimum vertex-cut graph partitioning algorithm and communication avoidance using a family of delayed-update algorithms. Our results on four common GNN benchmark datasets: Reddit, OGB-Products, OGB-Papers and Proteins, show up to 3.7x speed-up using a single CPU socket and up to 97x speed-up using 128 CPU sockets, respectively, over baseline DGL implementations running on a single CPU socket

研究动机与目标

• 解决在分布式 CPU 集群上对大规模图进行全批量 GNN 训练时的可扩展性挑战。
• 减少单节点上的内存和带宽瓶颈，并最小化分布式环境下的节点间通信。
• 通过共享内存和通信感知技术，优化深度图库（DGL）以实现在 CPU 集群上的高效全批量训练。
• 通过跨数百个计算节点的扩展，实现大规模图上 GNN 训练的可行性。

提出的方法

• 采用高效的共享内存实现，以减轻单个节点的内存压力。
• 应用最小顶点割图划分算法，以最小化节点间的通信量。
• 引入一系列延迟更新算法，以避免节点间的冗余通信。
• 通过集成通信减少和避免策略，优化 DGL 框架以支持分布式训练。
• 设计一种分布式训练流水线，在显著提升训练吞吐量的同时保持模型精度。

实验结果

研究问题

• RQ1 如何在数百个 CPU 节点上高效地扩展全批量 GNN 训练？
• RQ2 何种图划分策略能最小化分布式 GNN 训练中的通信量？
• RQ3 通过延迟更新实现的通信避免能在多大程度上提升训练性能？
• RQ4 所提出的框架与基线 DGL 实现相比，在加速比和可扩展性方面表现如何？

主要发现

• 在单个 CPU 插槽上，DistGNN 相较于基线 DGL 最多实现 3.7 倍的加速。
• 在 128 个 CPU 插槽上，DistGNN 相较于单插槽基线 DGL 训练最多实现 97 倍的加速。
• 最小顶点割图划分算法能有效减少分布式 GNN 训练中的节点间通信量。
• 延迟更新通信避免技术显著降低了节点间的通信开销。
• 该框架在保持模型性能的同时，实现了大规模图上可扩展的全批量训练。