A Dependency-Aware Multi-GPU Training Framework for Temporal Graph Neural Networks

Yang Zhao; Yue Dai; Liang Qiao; Buxin Tu; Bo Liu; Sen Zhang; Yi Zhang; Xiaoyu Hao

cs.LGOct 11, 2026

A Dependency-Aware Multi-GPU Training Framework for Temporal Graph Neural Networks

Authors: Yang Zhao, Yue Dai, Liang Qiao, Buxin Tu, Bo Liu, Sen Zhang, Yi Zhang, Xiaoyu Hao,

Organizations: University of Science and Technology of China, Hefei, Anhui, China · Illinois Institute of Technology, Chicago, Illinois, USA

Abstract

Temporal graph neural networks (TGNNs) achieve superior performance in many dynamic graph learning tasks. These models use node state vectors (i.e., memories) to capture and record graph changes (i.e., events), and then exploit these memories to assist prediction. However, this memory-centric design makes multi-GPU training challenging because predictions for later events often depend on memory updates from earlier events. To maintain model quality, existing multi-GPU TGNN training schemes strictly preserve all event dependencies, which limits parallelism and causes expensive cross-GPU memory synchronization. In this work, we observe that many of these dependencies can be handled more flexibly without harming training quality. Specifically, (a) dependencies in TGNN training have high locality in both time and graph topology, (b) the short-range dependencies missed within a large batch can often be recovered cheaply, and (c) many memory updates are not immediately needed. Based on these observations, we propose DepBloom to accelerate multi-GPU TGNN training with three techniques: (1) temporally bounded spatial partitioning to reduce unnecessary cross-GPU synchronization and improve parallelism; (2) intra-batch memory refinement to recover missing short-range temporal dependencies inside enlarged batches; and (3) critical-path optimization to prioritize urgent memory updates while deferring non-critical work during multi-GPU pipeline execution. Experimental results across representative models and diverse datasets show that DepBloom achieves an average speedup of 2.84 $\times$ and up to 5.16 $\times$ over state-of-the-art methods, while maintaining comparable or better model accuracy.

Figures & tables

Figure 1. TGNN computing steps in one batch. In many TGNN training pipelines, the memory update of one batch is deferred and executed as the first step of the next batch.

Figure 2. Comparison of four multi-GPU training paradigms for TGNNs: (a) batch parallelism, (b) epoch parallelism, (c) memory parallelism, and (d) pipeline parallelism. Here,

B_{j}^{i}

denotes the

j

-th batch in epoch

i

, and different colors distinguish batches from different epochs. RW denotes memory read and write operations, where different colors indicate different memory copies. Green arrows indicate inter-GPU data movement.

Figure 3. Latencies of training TGN on WIKI-TALK with different numbers of GPUs. (1) Sample&fetch includes neighbor sampling and input transfer from CPU to GPU; (2) memory_sync denotes the transfer of updated node memories across GPUs; (3) memory_update computes new node-memory states required by subsequent GPUs; and (4) model_train includes forward propagation, backward propagation, and parameter updates.

Figure 4. Trade-off between training efficiency and model quality under different batch sizes on WIKI-TALK .

Figure 5. (a) Cross-chain event ratio versus number of sub-chains, comparing naive sequential partitioning against topology-aware partitioning. (b) Cross-chain event ratio and training loss as a function of temporal window size.

Figure 6. (a) Execution modes under different parallel strategies. (b) Speedup versus training loss increase under larger batch sizes. Arrows show the percentage loss increase from BS=1200 to BS=4800.

Figure 7. Urgency mismatch in batch-produced memory updates. (a) Distribution of updated nodes immediately required by the next batch ( Hot Node ) and those that are not ( Cold Node ). (b) Breakdown of the critical path latency.

Figure 9. Comparison between the original pipeline parallelism and Topology-Aware Hierarchical Multi-Chain Parallelism. Using eight consecutive batches (batch size=2) as an example, the original design executes batches in a single pipeline and incurs frequent unnecessary synchronization. Our design partitions events into multiple topology-aware sub-chains and maps each sub-chain to a GPU group for parallel execution across sub-chains with minimal dependency loss.

Figure 10. Comparison of three multi-GPU pipeline execution strategies for TGNN training. (a) Pipelined Small Batches : each GPU executes one small batch per stage in strict sequential order. (b) Pipelined Large Batches (Local Parallelism) : micro-batches within each large batch are scheduled in parallel, each consuming the same stale memory from the previous batch. (c) Pipelined Large Batches (Local Refinement) : same scheduling as (b), but a refinement step (green) is inserted between micro-batches to propagate short-range dependencies locally before the next micro-batch executes. : actual scheduling dependencies that enforce serial execution order; : data usage relationships among micro-batches, which allow parallel execution ; : sample & fetch; : memory update; : memory refinement; : training & backward.

Figure 11. Comparisons between unified node update and priority-aware decoupled node update.

Figure 12. Overall performance comparison on NVIDIA A100 and RTX 4090 across four datasets and five TGNN models. The first two rows show end-to-end speedup over DistTGL 1 , and the last two rows show average precision (AP).

	Sample Message	Memory Update	Node Embedding
JODIE	most recent (n=1)	RNN	Identity
TGN	most recent (n=1)	GRU	GAT
APAN	most recent (n=10)	Transformer	Identity
TGAT	uniform (n=10)	Identity	2-layer GAT
DySAT	uniform (n=10)	GAT	RNN

Table 2. Details of TGNN models, with

d_{out}=100

Dataset	$\lvert V\rvert$	$\lvert E\rvert$	$\lvert d_{v}\rvert$	$\lvert d_{e}\rvert$
WIKI	9,227	157,474	172	100*
REDDIT	11,000	672,447	172	100*
WIKI-TALK	1,140,149	7,833,140	32*	100*
GDELT	16,682	191,290,882	186	413

Table 3. Dataset statistics.

\lvert d_{v}\rvert

and

\lvert d_{e}\rvert

show the dimensions of node features and edge features, respectively.

Figure 13. Scalability on WIKI-TALK and GDELT with APAN , under single- and multi-node GPU configurations. For multi-node, inter-node memory parallelism is combined with DepBloom’s intra-node dependency-aware execution.

Figure 14. Effect of temporally bounded spatial partitioning. (a) Average precision (AP). (b) End-to-end speedup.

Figure 15. Effect of intra-batch memory refinement. SB uses batch size 600, LB uses batch size 2400, and LB+IMR(4) splits each large batch into four chunks of size 600. (a) Normalized validation loss. (b) End-to-end speedup.

Figure 16. Effect of critical-path optimization on multi-GPU pipeline execution. Comparison of pipeline execution with ( W ) and without ( W/O ) the optimization.

Figure 17. Sensitivity of intra-batch refinement and temporal bound. (a) Sensitivity to refinement granularity under different batch sizes on WIKI and REDDIT . Top: Normalized training loss. Bottom: Speedup. (b) Sensitivity to the temporal bound on REDDIT and WIKI-TALK . Normalized training loss.

Figure 18. Latency breakdown of DepBloom across datasets and TGNN models, including preprocessing for temporally bounded spatial partitioning and extra computation from intra-batch memory refinement.

Temporal GNNs

Temporal graph neural networks model how node relationships, graph structure, and features change over time to support prediction on evolving networks.

See all 82 papers

FAST: A Holistic Framework for Optimizing Memory-I/O, Computation, and Sampling in Temporal GNN Training
Jul 6, 2026Yushu Cai, Qingrui Zhu, Lei Liu +3Temporal GNNs Graph Neural Networks
Xidian University Guangzhou, Guangdong, China · Hunan University Changsha, Hunan, China
AnchorRM: Parallel Intra-Batch State Reconstruction with Anchor-Based Residual Memory for Temporal Graph Neural Networks
Oct 20, 2026Haroon Gharwi, Kai Shu, Yue DaiTemporal GNNs Memory-Augmented Neural Networks
Department of Computer Science, Illinois Institute of Technology, Chicago, IL, USA · Department of Computer Science, Emory University, Atlanta, GA, USA
RAFT: Redundancy-Aware Input Filtering for Temporal Graph Neural Network Training
Oct 17, 2026Yue Dai, Sarah Zhou, Youtao ZhangTemporal GNNs
Department of Computer Science, Illinois Institute of Technology, Chicago, USA · BASIS Independent Silicon Valley, San Jose, CA, USA · Department of Computer Science, University of Pittsburgh, Pittsburgh, USA

Jul 6, 2026J/K move · Enter open · S save

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

Yushu Cai, Qingrui Zhu, Lei Liu, Kai Sheng+2

Xidian University Guangzhou, Guangdong, China · Hunan University Changsha, Hunan, China

Temporal GNNs Graph Neural Networks GPU Kernel Optimization

Open paper AskPDF arXiv

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.

Abstract

Temporal GNNs

Explore similar work

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

AnchorRM: Parallel Intra-Batch State Reconstruction with Anchor-Based Residual Memory for Temporal Graph Neural Networks

RAFT: Redundancy-Aware Input Filtering for Temporal Graph Neural Network Training