提出新框架，让时空图神经网络在超大规模数据上高效训练
PGT-I: Scaling Spatiotemporal GNNs with Memory-Efficient Distributed Training
- 用动态索引批处理技术减少内存占用
- 在128块GPU上实现11.78倍加速，峰值内存降89%
- 适合做大规模交通、气象等时空建模的研究者
时空图神经网络（ST-GNN）能有效建模空间与时间依赖关系，但受限于内存，仅适用于小规模数据。现有分布式训练框架不支持时空模型，且忽视其特性。基于对大规模任务的扩展性研究，我们提出 PyTorch Geometric Temporal Index（PGT-I），在 PyG Temporal 基础上集成分布式数据并行训练，并引入两种新策略：索引批处理和分布式索引批处理。这些方法利用时空结构在运行时动态构建快照，显著降低内存开销；分布式索引批处理进一步实现多GPU可扩展处理。该方法首次实现无需图划分即可在完整 PeMS 数据集上训练 ST-GNN，峰值内存降低最高达 89%，在 128 块 GPU 上相比标准 DDP 实现最高 11.78 倍加速。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。