Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Xin Huang, Weipeng Zhuo, Minh Phu Vuong, Shiju Li, Jongryool Kim, Bradley Rees, Chul-Ho Lee

TL;DR

SDT-GNN tackles the memory bottleneck in distributed GNN training on large graphs by streaming edges for partitioning and introducing SPRING to reduce replication. It uses weight-averaging synchronization to enable training when total GPU memory is insufficient, achieving up to 95% memory reduction while preserving accuracy across seven datasets. The framework integrates SPRING with flexible workload distribution and supports multiple GNN models, delivering comparable results to centralized training and enabling previously infeasible large-scale experiments (e.g., OGB-Papers). Experimental results demonstrate SPRING's efficiency gains and SDT-GNN's scalability on commodity hardware, offering a practical, memory-efficient path to billion-scale GNNs.

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.

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

TL;DR

SDT-GNN tackles the memory bottleneck in distributed GNN training on large graphs by streaming edges for partitioning and introducing SPRING to reduce replication. It uses weight-averaging synchronization to enable training when total GPU memory is insufficient, achieving up to 95% memory reduction while preserving accuracy across seven datasets. The framework integrates SPRING with flexible workload distribution and supports multiple GNN models, delivering comparable results to centralized training and enabling previously infeasible large-scale experiments (e.g., OGB-Papers). Experimental results demonstrate SPRING's efficiency gains and SDT-GNN's scalability on commodity hardware, offering a practical, memory-efficient path to billion-scale GNNs.

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.
Paper Structure (19 sections, 1 theorem, 11 figures, 5 tables)

This paper contains 19 sections, 1 theorem, 11 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Background and Motivations
  3. Distributed GNN Training
  4. Limitations in Distributed GNN Training
  5. SDT-GNN
  6. Overview
  7. Graph Partitioning
  8. Workload Allocation
  9. Distributed Model Training
  10. Implementation
  11. Evaluation
  12. Setup
  13. Streaming Partitioning Algorithms
  14. GNNs Trained Using SDT-GNN
  15. Comparison with DistDGL and PyG
  16. ...and 4 more sections

Key Result

Theorem 1

SPRING has a time complexity of $\mathcal{O}(|E|+|V|\log|V|)$ and a space complexity of $\mathcal{O}(|V|)$.

Figures (11)

  • Figure 1: End-to-end workflow of a distributed GNN training system.
  • Figure 2: Node classification accuracy of (a) GAT and (b) GraphSAGE trained on partitioned subgraphs by 2PS-L with (w/) and without (w/o) full neighbors added to the subgraphs.
  • Figure 3: Overview of SDT-GNN.
  • Figure 4: Illustration of SPRING as a part of SDT-GNN.
  • Figure 5: Training workflow of SDT-GNN after the input graph is partitioned.
  • ...and 6 more figures

Theorems & Definitions (2)

  • Theorem 1
  • proof