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

TL;DR

RapidGNN introduces a deterministic sampling strategy that precomputes mini-batches, enabling efficient cache and prefetching mechanisms to significantly reduce communication overhead and improve training throughput for large-scale GNNs.

Contribution

It proposes a novel deterministic sampling approach that anticipates feature access patterns, optimizing communication and caching in distributed GNN training.

Findings

Achieves 2.10x faster training throughput on average

Reduces remote feature fetches by over 4x

Cuts energy consumption by up to 23%

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%.