DistGNN: Scalable Distributed Training for Large-Scale Graph Neural Networks
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.7 $ speed-up using a single CPU socket and up to $97 $ speed-up using $128$ CPU sockets, respectively, over baseline DGL implementations running on a single CPU socket.
Benchmark Results
Tested on 3 datasets Non-standard metric on 1 dataset
Tables show presented results in global context (best 3 papers + GCN + MLP).
Classic
Sorted by Accuracy. Highlighted rows are from other papers.
| # | Model | Accuracy ↑ | Paper |
|---|---|---|---|
| 1 | GAP-EDP | 0.9870
±0.0003 | |
| 2 | NE-ASGCN | 0.9758
±0.0027 | |
| 3 | HSGT | 0.9730
±0.0024 | |
| 21 | GraphSAGE | 0.9653
±0.0011 | |
| 85 | cd-0 | 0.9370 | This paper |
| 87 | cd-5 | 0.9359 | This paper |
| 88 | 0c | 0.9358 | This paper |
| 147 | MLP | 0.8240
±0.0010 |
OGB
| # | Model | Accuracy ↑ | Paper |
|---|---|---|---|
| 1 | Jacobi | 0.8968 | |
| 2 | GPR | 0.8953 | |
| 3 | LEMP4HG | 0.8939 | |
| 4 | SAGE | 0.8933 | |
| 179 | cd-5 | 0.7918 | This paper |
| 181 | 0c | 0.7910 | This paper |
| 225 | cd-0 | 0.7763 | This paper |
| 351 | MLP | 0.6605
±0.0020 |
| # | Model | Accuracy ↑ | Paper |
|---|---|---|---|
| 1 | SGC | 0.7910 | |
| 2 | GAMLP | 0.7036
±0.0002 | |
| 3 | GLEM-GNN | 0.7036
±0.0002 | |
| 35 | SAGE | 0.6644
±0.0017 | |
| 95 | MLP | 0.4960
±0.0030 | |
| 99 | DistGNN | 0.4129 | This paper |
