DistGNN: Scalable Distributed Training for Large-Scale Graph Neural Networks
This addresses the challenge of training GNNs on large graphs for researchers and practitioners, though it is incremental as it optimizes an existing library.
The paper tackles the problem of scaling full-batch training for large-scale Graph Neural Networks (GNNs) on CPU clusters, achieving up to 3.7x speed-up on a single CPU socket and up to 97x speed-up using 128 CPU sockets over baseline implementations.
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