arrow
Return

NeutronTP: Load-Balanced Distributed Full-Graph GNN Training with Tensor Parallelism

delete2025-02-28
delete0
PRE
AI
X
Xin Ai *
H
Hao Yuan
L
Ling, Zeyu
Q
Qiange Wang
F
Fu, Zhenbo
C
Chen, Chaoyi
G
Gu, Yu
Y
Yu, Ge
DOI:10.14778/3705829.3705837delete
deleteOriginal
deleteOriginal request for help
deleteShare
deleteSave
Abstract

Abstract

En 中文
Graph neural networks (GNNs) have emerged as a promising direction. Training large-scale graphs that relies on distributed computing power poses new challenges. Existing distributed GNN systems leverage data parallelism by partitioning the input graph and distributing it to multiple workers. However, due to the irregular nature of the graph structure, existing distributed approaches suffer from unbalanced workloads and high overhead in managing cross-worker vertex dependencies. In this paper, we leverage tensor parallelism for distributed GNN training. GNN tensor parallelism eliminates cross-worker vertex dependencies by partitioning features instead of graph structures. Different workers are assigned training tasks on different feature slices with the same dimensional size, leading to a complete load balance. We achieve efficient GNN tensor parallelism through two critical functions. Firstly, we employ a generalized decoupled training framework to decouple NN operations from graph aggregation operations, significantly reducing the communication overhead caused by NN operations which must be computed using complete features. Secondly, we employ a memory-efficient task scheduling strategy to support the training of large graphs exceeding single GPU memory, while further improving performance by overlapping communication and computation. By integrating the above techniques, we propose a distributed GNN training system NeutronTP. Our experimental results on a 16-node Aliyun cluster demonstrate that NeutronTP achieves 1.29x-8.72x speedup over state-of-the-art GNN systems including DistDGL, NeutronStar, and Sancus.
Keywords:
NEURAL-NETWORKS

Journal

P
Proceedings of the VLDB Endowment
IF:
3.3
Papers:
556
Citations:
1.2W

Organization

No organization information available