FAST: A Holistic Framework for Optimizing Memory-I/O, Computation, and Sampling in Temporal GNN Training

2026-07-06Machine Learning

Machine Learning
AI summary

The authors address the challenge of training Temporal Graph Neural Networks (TGNNs) on large, changing graphs, which is slowed down by memory handling, irregular processing, and finding neighbors over time. They introduce FAST, a system that speeds up the whole training process by improving data movement, computation efficiency, and sampling together. FAST uses a new caching method called SlimCache, optimized graph operations, and a smart sampling strategy to reduce delays. Their tests show FAST can train TGNNs more than twice as fast as current methods without losing accuracy.

Temporal Graph Neural Networksdynamic graphsmemory I/Ograph samplingGPU cacheaggregationedge softmaxCPU cache localityhost-device data movement
Authors
Yushu Cai, Qingrui Zhu, Lei Liu, Kai Sheng, Hao Chen, Xin He
Abstract
Temporal Graph Neural Networks (TGNNs) are widely used for learning from dynamic graphs in applications such as recommendation, social network analysis, and traffic forecasting. However, scaling TGNN training to large dynamic graphs remains challenging due to three intertwined bottlenecks: memory I/O, irregular computation, and temporal neighbor sampling. Existing systems often optimize these stages in isolation, leaving substantial performance headroom on the table. We present FAST, a holistic framework that accelerates end-to-end TGNN training by jointly optimizing sampling, memory I/O, and computation. FAST introduces SlimCache, which exploits within-batch compression and cross-batch caching to reduce host-device data movement under limited GPU memory budgets. It further designs thread-efficient graph operators tailored to sparse temporal subgraphs, improving GPU cache locality and reducing the latency of aggregation and edge softmax. In addition, FAST employs a topology-aware sampling strategy that improves CPU cache locality and accelerates temporal neighbor sampling. Extensive experiments on real-world large dynamic graphs show that FAST achieves an average of 2.1x (up to 4.7x) speedup over state-of-the-art systems without sacrificing model accuracy.