RapidGNN: Communication Efficient Large-Scale Distributed Training of Graph Neural Networks

2025-05-16

AI Summary

To address high communication overhead and memory pressure in distributed training of large-scale Graph Neural Networks (GNNs), this paper proposes a co-optimization framework based on deterministic pre-sampling. It models mini-batch feature access patterns to enable precise remote feature caching and prefetching, and introduces— for the first time—a deterministic precomputation mechanism that preserves training stochasticity while significantly reducing the frequency and latency of remote data fetches. The method integrates deterministic graph sampling, feature access modeling, distributed cache scheduling, and prefetching strategies into a lightweight, efficient training framework. Evaluated on Reddit and OGBN-Products, it achieves up to 2.45× (average 2.10×) end-to-end training throughput improvement, reduces remote feature accesses by over 4×, and lowers energy consumption by up to 23%.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Search and Optimization: Distributed Search
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Web Mining and Content Analysis: Large pretrained models with web data
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Graph Neural Networks (GNNs) have achieved state-of-the-art (SOTA) performance in diverse domains. However, training GNNs on large-scale graphs poses significant challenges due to high memory demands and significant communication overhead in distributed settings. Traditional sampling-based approaches mitigate computation load to some extent but often fail to address communication inefficiencies inherent in distributed environments. This paper presents RapidGNN that introduces a deterministic sampling strategy to precompute mini-batches. By leveraging the sampling strategy, RapidGNN accurately anticipates feature access patterns, enabling optimal cache construction and timely prefetching of remote features. This reduces the frequency and latency of remote data transfers without compromising the stochastic nature of training. Evaluations on Reddit and OGBN-Products datasets demonstrate that RapidGNN achieves significant reductions in training time and remote feature fetches, outperforming existing models in both communication efficiency and throughput. Our findings highlight RapidGNN's potential for scalable, high-performance GNN training across large, real-world graph datasets along with improving energy efficiency. Our model improves end-to-end training throughput by 2.10x on average over SOTA model GraphSAGE-METIS (up to 2.45x in some settings), while cutting remote feature fetches by over 4x. It also reduces energy consumption up to 23%.

Problem

• Reducing communication overhead in distributed GNN training
• Optimizing feature access patterns for efficient caching
• Improving scalability and energy efficiency in large-scale GNNs

Innovation

• Deterministic sampling for mini-batch precomputation
• Optimal cache construction via feature access prediction
• Reduced remote data transfers with timely prefetching