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

arXiv: https://arxiv.org/abs/2507.11683 · HTML: https://mlsanity.com/paper/2507.11683/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 78 references that are in the map.
- First posted on arXiv 15 Jul 2025.

## Builds on (most cited first, 30 of 34)

- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Spatio-Temporal Graph Convolutional Networks: A Deep Learning Framework for Traffic Forecasting](https://mlsanity.com/paper/1709.04875/) (2017)
- [Graph WaveNet for Deep Spatial-Temporal Graph Modeling](https://mlsanity.com/paper/1906.00121/) (2019)
- [Large Batch Training of Convolutional Networks](https://mlsanity.com/paper/1708.03888/) (2017)
- [Adaptive Graph Convolutional Recurrent Network for Traffic Forecasting](https://mlsanity.com/paper/2007.02842/) (2020)
- [T-GCN: A Temporal Graph ConvolutionalNetwork for Traffic Prediction](https://mlsanity.com/paper/1811.05320/) (2018)
- [One weird trick for parallelizing convolutional neural networks](https://mlsanity.com/paper/1404.5997/) (2014)
- [Dynamic Graph Convolutional Recurrent Network for Traffic Prediction: Benchmark and Solution](https://mlsanity.com/paper/2104.14917/) (2021)
- [DistDGL: Distributed Graph Neural Network Training for Billion-Scale Graphs](https://mlsanity.com/paper/2010.05337/) (2020)
- [PyTorch Geometric Temporal: Spatiotemporal Signal Processing with Neural Machine Learning Models](https://mlsanity.com/paper/2104.07788/) (2021)
- [PyTorch Distributed: Experiences on Accelerating Data Parallel Training](https://mlsanity.com/paper/2006.15704/) (2020)
- [Normalization Techniques in Training DNNs: Methodology, Analysis and Application](https://mlsanity.com/paper/2009.12836/) (2020)
- [Spatial-Temporal Large Language Model for Traffic Prediction](https://mlsanity.com/paper/2401.10134/) (2024)
- [A3T-GCN: Attention Temporal Graph Convolutional Network for Traffic Forecasting](https://mlsanity.com/paper/2006.11583/) (2020)
- [Learning the Evolutionary and Multi-scale Graph Structure for Multivariate Time Series Forecasting](https://mlsanity.com/paper/2206.13816/) (2022)
- [Distributed Graph Neural Network Training: A Survey](https://mlsanity.com/paper/2211.00216/) (2022)
- [GADMM: Fast and Communication Efficient Framework for Distributed Machine Learning](https://mlsanity.com/paper/1909.00047/) (2019)
- [Conditional Local Convolution for Spatio-temporal Meteorological Forecasting](https://mlsanity.com/paper/2101.01000/) (2021)
- [Parallel and Distributed Graph Neural Networks: An In-Depth Concurrency Analysis](https://mlsanity.com/paper/2205.09702/) (2022)
- [AST-GCN: Attribute-Augmented Spatiotemporal Graph Convolutional Network for Traffic Forecasting](https://mlsanity.com/paper/2011.11004/) (2020)
- [Learning Internal Representations (COLT 1995)](https://mlsanity.com/paper/1911.05781/) (2019)
- [Convergence Analysis of Distributed Stochastic Gradient Descent with Shuffling](https://mlsanity.com/paper/1709.10432/) (2017)
- [Chickenpox Cases in Hungary: a Benchmark Dataset for Spatiotemporal Signal Processing with Graph Neural Networks](https://mlsanity.com/paper/2102.08100/) (2021)
- [Graph-Partitioning-Based Diffusion Convolutional Recurrent Neural Network for Large-Scale Traffic Forecasting](https://mlsanity.com/paper/1909.11197/) (2019)
- [STGC-GNNs: A GNN-based traffic prediction framework with a spatial-temporal Granger causality graph](https://mlsanity.com/paper/2210.16789/) (2022)
- [DistTGL: Distributed Memory-Based Temporal Graph Neural Network Training](https://mlsanity.com/paper/2307.07649/) (2023)

## Built on by (most cited first, 3 of 3)

- [A Distributed Framework for Causal Modeling of Performance Variability in GPU Traces](https://mlsanity.com/paper/2510.18300/) (2025)
- [Learning from Complexity: Exploring Dynamic Sample Pruning of Spatio-Temporal Training](https://mlsanity.com/paper/2602.19113/) (2026)
- [Two-level domain-decomposition AdaGrad method for scalable training of graph neural networks](https://mlsanity.com/paper/2608.22575/) (2026)

## Abstract (arXiv)

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.