PGT-I: Scaling Spatiotemporal GNNs with Memory-Efficient Distributed Training

TL;DR

This paper introduces PGT-I, a scalable distributed training framework for spatiotemporal GNNs that significantly reduces memory usage and accelerates training on large datasets by exploiting spatiotemporal data properties.

Contribution

The paper presents PGT-I, a novel extension to PyTorch Geometric Temporal, with index-batching strategies enabling scalable, memory-efficient training of spatiotemporal GNNs across multiple GPUs.

Findings

Reduced peak memory usage by up to 89%.

Achieved up to 11.78x speedup with 128 GPUs.

Enabled training on entire PeMS dataset without graph partitioning.

Abstract

Spatiotemporal graph neural networks (ST-GNNs) are powerful tools for modeling spatial and temporal data dependencies. However, their applications have been limited primarily to small-scale datasets because of memory constraints. While distributed training offers a solution, current frameworks lack support for spatiotemporal models and overlook the properties of spatiotemporal data. Informed by a scaling study on a large-scale workload, we present PyTorch Geometric Temporal Index (PGT-I), an extension to PyTorch Geometric Temporal that integrates distributed data parallel training and two novel strategies: index-batching and distributed-index-batching. Our index techniques exploit spatiotemporal structure to construct snapshots dynamically at runtime, significantly reducing memory overhead, while distributed-index-batching extends this approach by enabling scalable processing across multiple GPUs. Our techniques enable the first-ever training of an ST-GNN on the entire PeMS dataset without graph partitioning, reducing peak memory usage by up to 89% and achieving up to a 11.78x speedup over standard DDP with 128 GPUs.