Plexus: Taming Billion-edge Graphs with 3D Parallel Full-graph GNN Training

TL;DR

This work tackles the challenge of training Graph Convolutional Networks on billion-edge graphs without sampling, a problem severely limited by memory, communication, and load balance at scale. It introduces Plexus, a three-dimensional tensor-parallel framework that distributes all matrices across a 3D GPU grid and cycles adjacency shards across layers to enable full-graph training on thousands of GPUs. A dedicated performance model predicts optimal 3D configurations, while optimizations such as a double permutation scheme and blocked aggregation improve load balancing and communication efficiency. Plexus demonstrates state-of-the-art scalability, achieving up to 2048 GPUs on Perlmutter and 1024 GPUs on Frontier with substantial speedups over prior frameworks, signaling a practical path to training full-graph GNNs on massive graphs.

Abstract

Graph neural networks (GNNs) leverage the connectivity and structure of real-world graphs to learn intricate properties and relationships between nodes. Many real-world graphs exceed the memory capacity of a GPU due to their sheer size, and training GNNs on such graphs requires techniques such as mini-batch sampling to scale. The alternative approach of distributed full-graph training suffers from high communication overheads and load imbalance due to the irregular structure of graphs. We propose a three-dimensional (3D) parallel approach for full-graph training that tackles these issues and scales to billion-edge graphs. In addition, we introduce optimizations such as a double permutation scheme for load balancing, and a performance model to predict the optimal 3D configuration of our parallel implementation -- Plexus. We evaluate Plexus on six different graph datasets and show scaling results on up to 2048 GPUs of Perlmutter, and 1024 GPUs of Frontier. Plexus achieves unprecedented speedups of 2.3-12.5x over prior state of the art, and a reduction in time-to-solution by 5.2-8.7x on Perlmutter and 7.0-54.2x on Frontier.

Figures (10)

• Figure 1: Different paradigms of GNN training that can be combined together, shown in four quadrants. Each quadrant shows a sample graph and its adjacency matrix. Blue nodes are part of the batch and grey nodes are not. Solid lines indicate edges considered during aggregation, and dashed line represent edges that are not considered. Red values in the adjacency matrix indicate that an entry has been modified.
• Figure 2: An overview of the 3D tensor parallel algorithm for GNN training. Eight GPUs are arranged in a 3D grid (X=Y=Z=2) and matrices in layer 0 of the network are distributed across different planes (shown in different colors).
• Figure 3: Shapes of the matrix shards (sub-blocks) in the first layer on a single GPU, showing two key matrix multiplications in the forward pass.
• Figure 4: Applying the 3D tensor parallel algorithm to all layers of a 3-layer GCN, connecting the output of one layer to the input of the next using unique shards of the adjacency matrix.
• Figure 5: Validating the performance model for the ogbn-products dataset on 64 GPUs of Perlmutter.