ATLAS: Enabling an Asynchronous State-Aware Pipeline for Efficient Training of Temporal Graph Neural Networks

Abstract

Temporal Graph Neural Networks (T-GNNs) are widely used to model continuous-time dynamic graphs. However, training T-GNNs on large temporal graphs remains computationally prohibitive due to a fundamental challenge: existing systems enforce a stop-line execution paradigm that serializes batch data preparation with GPU computation, leaving host-side preparation, data movement, and fragmented GPU submission on the critical path. We present Atlas, a principled training system that addresses this bottleneck through a complementary asynchronous state-aware pipeline design. We decouple bulk synchronization from temporal dependency by introducing a temporal state model that formally decouples history-bound batch preparation from state-consuming execution, enabling a state-versioned inter-batch pipeline that overlaps host-memory structure preparation with GPU execution while preserving dependent training semantics. This design enables a new opportunity for fully asynchronous parallelism in T-GNN training. Further, Atlas employs template-stabilized GPU execution that converts irregular temporal neighborhoods into a bounded set of fixed-shape tensor templates, enabling CUDA Graph capture and replay to amortize kernel launch overhead. Across four models and four standard datasets, Atlas outperforms TGLite in all 16 configurations, delivering geometric-mean speedups of 4.70 × on H100 and 5.55 × on RTX PRO 6000, with a peak speedup of 11.60 ×. Across the TGAT and TGN configurations supported by Pulse, Atlas provides a further 1.12 × –3.88 × speedup on the two platforms. Atlas also consistently reduces capped training-window time on large WikiTalk and GDELT stress-test slices. Profiling shows that Atlas reduces kernel launches by \(81.5\%\)–\(99.7\%\) while preserving predictive accuracy across the standard benchmarks. The source code is publicly available at https://github.com/yongqiu-star/ATLAS.