DistGNN: Scalable Distributed Training for Large-Scale Graph Neural Networks

Vasimuddin, Sanchit Misra, Guixiang Ma, Ramanarayan Mohanty, E. Georganas, A. Heinecke, Dhiraj D. Kalamkar, Nesreen Ahmed, Sasikanth Avancha

Main figure for 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.

Proposed method: DistGNN

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

Reddit

Global leaderboard →
Non-default metric. This paper reports Accuracy; the dataset default is F1. Values below use Accuracy only.

Sorted by Accuracy. Highlighted rows are from other papers.

#ModelAccuracy ↑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

#ModelAccuracy ↑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
SOTA performance (not significantly different from the best) Proposed method (this paper) Baseline / comparison (this paper) Global context (other papers)