SDT-GNN: Streaming-based Distributed Training Framework for Graph Neural Networks

TL;DR

SDT-GNN introduces a streaming-based distributed training framework for GNNs that significantly reduces memory usage and enables training on large graphs with limited GPU memory, outperforming existing methods.

Contribution

It proposes SDT-GNN, a novel streaming-based framework with SPRING, a new partitioning algorithm, allowing efficient GNN training on large graphs with lower memory requirements.

Findings

SDT-GNN reduces memory footprint by up to 95% compared to existing frameworks.

SPRING outperforms state-of-the-art streaming partitioning algorithms.

SDT-GNN maintains prediction accuracy while using less memory.

Abstract

Recently, distributed GNN training frameworks, such as DistDGL and PyG, have been developed to enable training GNN models on large graphs by leveraging multiple GPUs in a distributed manner. Despite these advances, their memory requirements are still excessively high, thereby hindering GNN training on large graphs using commodity workstations. In this paper, we propose SDT-GNN, a streaming-based distributed GNN training framework. Unlike the existing frameworks that load the entire graph in memory, it takes a stream of edges as input for graph partitioning to reduce the memory requirement for partitioning. It also enables distributed GNN training even when the aggregated memory size of GPUs is smaller than the size of the graph and feature data. Furthermore, to improve the quality of partitioning, we propose SPRING, a novel streaming partitioning algorithm for distributed GNN training. We demonstrate the effectiveness and efficiency of SDT-GNN on seven large public datasets. SDT-GNN has up to 95% less memory footprint than DistDGL and PyG without sacrificing the prediction accuracy. SPRING also outperforms state-of-the-art streaming partitioning algorithms significantly.