Simplifying Distributed Neural Network Training on Massive Graphs: Randomized Partitions Improve Model Aggregation

Jiong Zhu, Aishwarya Reganti, Edward Huang, Charles Dickens, Nikhil Rao, Karthik Subbian, Danai Koutra

Main figure for Simplifying Distributed Neural Network Training on Massive Graphs: Randomized Partitions Improve Model Aggregation

Abstract

Distributed training of GNNs enables learning on massive graphs (e.g., social and e-commerce networks) that exceed the storage and computational capacity of a single machine. To reach performance comparable to centralized training, distributed frameworks focus on maximally recovering cross-instance node dependencies with either communication across instances or periodic fallback to centralized training, which create overhead and limit the framework scalability. In this work, we present a simplified framework for distributed GNN training that does not rely on the aforementioned costly operations, and has improved scalability, convergence speed and performance over the state-of-the-art approaches. Specifically, our framework (1) assembles independent trainers, each of which asynchronously learns a local model on locally-available parts of the training graph, and (2) only conducts periodic (time-based) model aggregation to synchronize the local models. Backed by our theoretical analysis, instead of maximizing the recovery of cross-instance node dependencies---which has been considered the key behind closing the performance gap between model aggregation and centralized training---, our framework leverages randomized assignment of nodes or super-nodes (i.e., collections of original nodes) to partition the training graph such that it improves data uniformity and minimizes the discrepancy of gradient and loss function across instances. In our experiments on social and e-commerce networks with up to 1.3 billion edges, our proposed RandomTMA and SuperTMA approaches---despite using less training data---achieve state-of-the-art performance and 2.31x speedup compared to the fastest baseline, and show better robustness to trainer failures.

Benchmark Results

Tested on 1 dataset

Tables show presented results in global context (best 3 papers + GCN + MLP).

OGB

#ModelMRR ↑Paper
1
CFG
0.8997 ±0.0015
2
SIEG
0.8987 ±0.0018
3
LPFormer
0.8981 ±0.0013
12
GCN
0.8705 ±0.0004
25
SuperTMA
0.8375 ±0.0043
This paper
28
RandomTMA
0.8328 ±0.0024
This paper
34
PSGD-PA
0.8240 ±0.0028
This paper
36
GGS
0.8195 ±0.0020
This paper
37
LLCG
0.8188 ±0.0002
This paper
71
MLP
0.3917 ±0.0044
SOTA performance (not significantly different from the best) Baseline / comparison (this paper) Global context (other papers)