Source-linked AI summary
DistGNN: Scalable Distributed Training for Large-Scale Graph Neural Networks
Vasimuddin Md, Sanchit Misra, Guixiang Ma, Ramanarayan Mohanty, Evangelos Georganas, Alexander Heinecke, Dhiraj Kalamkar, Nesreen K. Ahmed, Sasikanth Avancha
TL;DR
Large-scale full-batch GNN training is constrained by memory, bandwidth, and communication demands. DistGNN optimizes CPU-cluster training through single-socket improvements, vertex-cut partitioning, and delayed updates, achieving reported speedups over DGL baselines.
Problem
Large-graph full-batch GNN training requires high memory capacity, memory bandwidth, and communication volume, limiting training on a single socket.
Method
DistGNN combines single-socket aggregation optimizations with vertex-cut partitioning and delayed partial aggregates to reduce and overlap communication.
Results
3.66× end-to-end training speedup was achieved over DGL 0.5.3 on workloads fitting on one socket, while Reddit reached 3.1× over the reported baseline.
Takeaways & Limitations
DistGNN demonstrates scalable full-batch GNN training across CPU sockets by mitigating aggregation and communication bottlenecks on large benchmark graphs.
Takeaways & Limitations
Proteins training accuracy could not be validated, and further hyperparameter tuning is required for OGBN-Papers.
Abstract
from arXiv · showhide
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.7x speed-up using a single CPU socket and up to 97x speed-up using 128 CPU sockets, respectively, over baseline DGL implementations running on a single CPU socket
1 INTRODUCTION
DistGNN targets the scalability barriers of full-batch GNN training on large graphs, including memory, bandwidth, and communication demands. It combines single-socket optimizations with distributed techniques to scale training across CPU clusters.
- Motivation: Large-graph GNN training is constrained by memory capacity, bandwidth requirements, and communication volume across compute nodes.These constraints limit the problem size solvable on one socket and can make communication a bottleneck.
- Motivation: Full-batch training can provide higher accuracy than neighborhood sampling in some cases, motivating distributed full-batch solutions.The paper focuses on full-batch training across Intel Xeon CPU sockets rather than mini-batch scaling.
- Motivation: Existing GNN libraries couple deep-learning primitives with graph message passing but can be inefficient for single-socket shared-memory execution.This motivates optimizing DGL’s core compute primitives for CPU systems.
- Present Work: DistGNN combines architecture-optimized aggregation with distributed solutions for large-scale GNN training on Intel Xeon CPU clusters.Its contributions include shared-memory optimization, vertex-cut partitioning, and delayed-update communication avoidance.
- Present Work: 97× speedup was achieved for GraphSAGE full-batch training on OGB-Papers across 128 Intel Xeon CPU sockets versus an unoptimized DGL baseline.The OGB-Papers dataset contains 111 million vertices and 1.6 billion edges.
2 BACKGROUND
GNNs learn graph structure by aggregating neighborhood features into vertex or edge embeddings, with DGL expressing these operations through message passing and matrix-multiplication primitives.
- GNN Aggregation: GNN aggregation recursively gathers multi-hop neighborhood features to encode vertex or edge representations.The aggregation function also learns shared weights through a shallow neural network.
- Aggregation Primitive: The aggregation primitive takes vertex and edge features, element-wise operators, and a reducer to produce output features.Vertex and edge feature matrices have sizes |V|×d and |E|×d, respectively.
- Aggregation Primitive: Unary aggregation copies one input feature stream before reducing it into the output.This is the unary form used when one operator input is NULL.
- DGL: DGL formulates aggregation as SpMM, f_O = A × f_X, using the graph adjacency matrix and dense feature matrix.DistGNN later develops SpMM variants and architecture-aware optimizations for these computations.
- DGL: DGL represents vertex message passing with SpMM and edge message passing with sampled dense-dense matrix multiplication.It supports built-in and user-defined binary and reduction operators.
3 RELATED WORK
Prior work covers distributed GNN training on GPUs and CPUs, but CPU-based approaches had focused on mini-batch training, leaving distributed full-batch CPU scaling comparatively unaddressed.
- Distributed GNN Training: Distributed GNN systems commonly partition input graphs across CPU sockets or GPU cards because model sizes are relatively small.The literature includes both full-batch and mini-batch approaches.
- GPU Systems: GPU systems such as NeuGraph, Roc, and CAGNET use specialized graph abstractions, partitioning, regression, or matrix-blocking strategies for distributed training.These systems demonstrate several approaches to distributed GNN computation on GPUs.
- CPU Systems: CPU-based prior approaches supported mini-batch training, distributed storage, sampling, aggregation, and graph partitioning, but did not report scaling for their solution.AliGraph also explored caching neighbors of important vertices to reduce communication.
- Graph Processing Frameworks: Graph-processing frameworks based on synchronous Gather-Apply-Scatter can implement aggregation but lack support for varied GNN operations and graph attention models.This limits their direct coverage of modern GNN workloads.
4 DISTGNN: SHARED-MEMORY ALGORITHM
DistGNN accelerates DGL’s aggregation primitive on shared-memory CPUs by addressing irregular memory access, load imbalance, and vectorization. Its implementation combines cache blocking, dynamic scheduling, and LIBXSMM-based loop optimization.
- Baseline and Optimization: Aggregation is a dominant runtime component in GNN applications, motivating optimized implementations within DGL.The shared-memory approach focuses on accelerating this primitive.
- Baseline Implementation: DGL’s aggregation primitive pulls source vertex and edge messages and reduces them into destination vertex outputs.The computation applies binary or unary operators followed by an element-wise reduction.
- Baseline Implementation: The baseline parallelizes destination vertices across OpenMP threads so each output feature vector has a single owning thread.This avoids write collisions during reduction.
- Performance Challenges: Large sparse graphs create streaming edge accesses but irregular, repeated source-feature gathers that stress memory bandwidth and cache locality.Feature matrices may not fit in cache, and neighbor features can be non-contiguously located.
- Cache Blocking: Cache blocking partitions source features so they can be reused from cache while processing destination vertices.The method reads source features once per block but makes multiple passes over output features.
- Multithreading: Dynamic OpenMP scheduling handles power-law variation in vertex degree while assigning contiguous destination ranges to threads.This supports balanced work distribution and contiguous output writes.
- Vectorization: LIBXSMM reorders loops and supplies architecture-optimized primitives for vectorizing the aggregation computation.It reduces the need for manually written SIMD intrinsics across DGL’s operator variants.
5 DISTGNN: DISTRIBUTED-MEMORY ALGORITHM
DistGNN distributes full-batch GraphSAGE training by partitioning the graph across CPU sockets while replicating the smaller model. Its algorithms trade communication, accuracy, and performance through local aggregation, synchronized partial aggregates, and delayed communication.
- Distributed-memory design: DistGNN uses data parallelism: the model is replicated across sockets while the input graph is partitioned.The approach targets aggregation, whose memory demands grow with graph size.
- Graph partitioning: Communication is minimized by partitioning the graph, although fully self-contained partitions would reduce communication at the cost of training accuracy.The design therefore treats communication reduction and accuracy as a practical trade-off.
- Graph partitioning: Vertex-cut partitioning distributes edges across partitions, allowing vertices to appear in multiple partitions with partial neighborhoods.Each split vertex carries a feature-vector copy and participates in local aggregation.
- Distributed aggregation algorithms: 0c ignores remote partial neighborhoods, completely avoids communication, and provides the fastest performance roofline across sockets.The algorithm is also the most optimistic with respect to accuracy.
- Distributed aggregation algorithms: cd-0 communicates partial aggregates to remote clones and waits for completion, so each partition receives its complete neighborhood.This is expected to match the single-socket algorithm’s accuracy while providing a lower-bound performance reference.
- Distributed aggregation algorithms: cd-r overlaps remote aggregate communication with local computation by delaying updates across epochs and optionally binning split-vertices.The reported accuracy is within 1% of the state-of-the-art baseline for each dataset.
6 EXPERIMENTAL EVALUATION
Experiments evaluate DistGNN on CPU hardware, benchmark datasets, single-socket optimizations, distributed scaling, accuracy, and workload differences from Dist-DGL. Results show that memory-aware optimizations and partitioning/delayed communication improve performance, while partition quality and dataset coverage constrain outcomes.
- Experiment Setup: Experiments use Intel Xeon CPUs, DGL 0.5.3, PyTorch 1.6.0, and five graph datasets, with reported epoch times averaged across specified epoch ranges.Single-socket runs use Xeon 8280 CPUs with 98 GB memory; distributed runs use Xeon 9242 systems connected by Mellanox HDR interconnect.
- Single-Socket Performance: 3.66× end-to-end speedup and 4.41× AP speedup are achieved over DGL 0.5.3 on workloads fitting within a 98 GB single socket.The optimized implementation confirms that aggregation often dominates application runtime.
- Single-Socket Performance: Blocking balances feature-vector reuse against additional output passes, with the best performance at the point minimizing total data read and written.Denser Reddit graphs place this reuse sweet spot at a larger block count than sparser OGBN-Products graphs.
- Single-Socket Performance: Dynamic scheduling primarily benefits OGBN-Products, blocking strongly benefits Reddit, and LIBXSMM loop reordering and JITing improve both datasets.The optimization effects correlate with total memory IO and execution time.
- Distributed Algorithm Performance: 123.13× speedup is reached for OGBN-Papers on 128 sockets with 0c, while OGBN-Products reaches 16.1× on 64 sockets with 0c.For OGBN-Papers, cd-0 and cd-5 reach 27.43× and 83.16×, respectively; replication quality and memory placement affect scaling.
- Distributed Algorithm Performance: Distributed algorithms remain within 1% of best accuracy on Reddit and OGBN-Products, but delayed updates with r=10 degrade accuracy because aggregates become increasingly stale.The experiments use r=5 for cd-r, while smaller delays showed no discernible accuracy improvement.
- Limitations: The evaluation cannot validate Proteins training accuracy and reports substantially lower vanilla GraphSAGE accuracy than the OGB leaderboard for OGBN-Papers.The authors identify missing Proteins features and labels and require further OGBN-Papers hyperparameter tuning.
- Distributed Algorithm Performance: DistGNN performs approximately 4×–13× and 13× more work per epoch than Dist-DGL on OGBN-Products yet reports comparable or better epoch time on similar hardware.The higher work reflects complete-neighborhood aggregation rather than the sampled workload used by Dist-DGL.
7 CONCLUSION AND FUTURE WORK
DistGNN provides a distributed full-batch GNN training solution that combines optimized aggregation with communication reduction and delayed partial aggregates. Its delayed partial aggregate algorithms generally preserve accuracy as partitions increase, while future work targets mini-batch training, broader models, and lower-precision communication.
- 7 CONCLUSION AND FUTURE WORK: DistGNN targets full-batch GNN training on Intel Xeon CPU clusters by optimizing aggregation and mitigating communication bottlenecks.The approach uses architecture-aware aggregation optimizations, vertex-cut partitioning, and communication overlap across epochs.
- 7 CONCLUSION AND FUTURE WORK: Future work will extend DistGNN to scalable mini-batch training, additional GNN models, and low-precision formats such as FP16 and BFLOAT16.The authors also plan broader accuracy analyses for 0c and cd-r across datasets and model architectures.