Spatio-temporal graph neural networks (ST-GNNs), which are effective for modeling large-scale spatio-temporal data dependencies, have been mainly applied to small-scale datasets due to memory constraints. In this paper, we present the PyTorch Geometric Temporal Index (PGT-I), an extension of PyTorch Geometric Temporal that integrates distributed data parallel learning and two novel strategies: index placement and distributed index placement. The indexing technique leverages spatio-temporal structures to dynamically generate snapshots at runtime, which significantly reduces memory overhead, while distributed index placement enables scalable processing across multiple GPUs. The proposed technique enables ST-GNNs to be trained on the entire PeMS dataset for the first time without graph partitioning, achieving up to 89% peak memory usage reduction and up to 11.78x speedup over standard DDP using 128 GPUs.
