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

The gist: DepBloom accelerates multi-GPU TGNN training by introducing three techniques—temporally bounded spatial partitioning, intra-batch memory refinement, and critical-path optimization—to handle dependencies more flexibly without harming training quality

How it works

DepBloom proposes a framework that addresses the challenges of scaling multi-GPU TGNN training by observing that many dependencies can be handled more flexibly than existing strict preservation schemes. The framework integrates three key techniques to increase parallelism and reduce cross-GPU synchronization costs.

The three core techniques are:

1. Temporally bounded spatial partitioning, which divides the event stream into serial temporal windows and organizes each window into topology-aware parallel sub-chains, thereby reducing unnecessary cross-GPU synchronization among sub-chains while preserving critical dependencies among windows.

2. Intra-batch memory refinement, which recovers missing short-range temporal information inside enlarged batches by splitting a large batch into consecutive micro-batches and using a lightweight refinement module to approximate the effect of omitted intermediate updates without requiring cross-device synchronization.

3. Critical-path optimization for multi-GPU pipeline execution, which prioritizes urgent memory updates and removes non-critical updates and operations from the critical path, thereby improving pipeline execution efficiency and enhancing system scalability.

Key Observations Driving the Design

The paper identifies three key observations that motivate the design of DepBloom.

** Dependencies in TGNN training have high locality in both time and graph topology, which is hard to utilize with only one-dimensional partitioning. Temporal-only partitioning breaks topological dependencies while spatial-only partitioning degrades training quality over long horizons. **

** The short-range temporal dependencies missed inside a large batch can often be recovered cheaply, specifically by using node memories carried over from the previous batch along with the most recent messages from earlier micro-batches to refine memory, enabling parallel micro-batch execution without synchronization. **

** Memory updates are not equally urgent for subsequent computation, as only a small subset of updated node memories is immediately required by the next batch, allowing non-urgent updates to be removed from the critical path. This suggests that current synchronization policies are overly conservative and place non-urgent updates on the critical path. 4.4 Design of DepBloom outlines how these observations translate into the framework, combining temporally bounded spatial partitioning, intra-batch memory refinement, and critical-path optimization. 4.2 Temporally Bounded Spatial Partitioning details the hierarchical organization where the event stream is divided into temporal windows and then each window is partitioned along its spatial dimension into sub-chains. 4.3 Intra-batch Memory Refinement describes splitting a large batch into an ordered sequence of micro-batches and computing refined memory using the formula s(i)v = COMB UPDT s−v, m(i)v, REFINE E(<i)v, ti. 4.4 Critical-path optimization involves two aspects: removing non-urgent memory updates from the critical path via Priority-aware memory-update decoupling and moving repetitive preprocessing work out of latency-sensitive stages via Critical-path precomputation and overlap. 8. DepBloom's design overview illustrates how these components work together to improve multi-GPU TGNN training from top to bottom. 9. Comparison of three multi-GPU pipeline execution strategies for TGNN training shows that DepBloom's Pipelined Large Batches (Locally Refinement) strategy is superior, as it includes a refinement step between micro-batches to propagate short-range dependencies locally before the next micro-batch executes. 12. Figure 10 compares three multi-GPU pipeline execution strategies for TGNN training, showing that Pipelined Large Batches (Locally Refinement) is the most effective strategy. 4.4 Critical-path optimization for multi-GPU pipeline execution shows that removing unnecessary operations from the pipeline critical path directly shortens dependency-induced stalls, with average speedups of 1.11×, 1.22×, and 1.34× on WIKI, REDDIT, and WIKI-TALK respectively. 5. Evaluation shows that DepBloom achieves an average speedup of 2.84× and up to 5.16× over state-of-the-art methods while maintaining comparable or better model quality. 9. Cross-platform trends indicate that DepBloom delivers consistent gains across both hardware platforms, achieving average speedups of 2.97× and 2.70× across all baseline comparisons. The quality gains mainly come from DepBloom’s dependency-aware execution: temporally bounded spatial partitioning keeps strongly coupled events within the same sub-chain, while intra-batch memory refinement recovers recent dependencies lost in enlarged batches. 14. Figure 14 shows that with our optimization enabled, the average AP increases from 91.24% to 93.02%, together with a 1.19× speedup over the baseline. Figure 15 demonstrates that intra-batch memory refinement recovers most of the lost accuracy with little overhead, as normalized loss drops to 1.03 with LB+IMR(4). Figure 2 shows that DepBloom consistently outperforms all baselines in throughput, achieving the highest speedup, ranging from 1.15× to 5.16×, with an average of 2.84×. The results show that much of the lost short-range temporal signal in large batches can be recovered without exact sequential memory updates. 5.4 Optimization Analysis confirms the effectiveness of each optimization, showing that Temporally Bounded Spatial Partitioning jointly exploiting temporal and topological locality better balances accuracy-efficiency tradeoff. The impact of intra-batch memory refinement shows that IMR recovers the missing information while still avoiding costly sequential updates. The critical-path optimization for multi-GPU pipeline execution consistently improves pipeline efficiency, with average speedups of 1.11×, 1.22×, and 1.34× on WIKI, REDDIT, and WIKI-TALK respectively. The low preprocessing overhead arises because temporal window construction and topology-aware spatial partitioning run before training and are reused during execution. 4.4 Critical-path optimization for multi-GPU pipeline execution shows that removing unnecessary operations from the pipeline critical path directly shortens dependency-induced stalls. 7. Conclusion summarizes that DepBloom combines temporally bounded spatial partitioning, intra-batch memory refinement, and critical-path optimization to reduce unnecessary synchronization, recover missing short-range temporal information, and shorten memory-synchronization stalls. 15. The paper concludes that DepBloom achieves 2.84× average speedup and up to 5.16× speedup over state-of-the-art baselines while maintaining comparable or better model quality. The final sentence states that the consistent gains across both CTDG and DTDG models demonstrate the generality of DepBloom across different temporal modeling paradigms. 48. The references section lists numerous prior works, including those on DistTGL [55] and PipeTGL [26], which are compared against DepBloom. 49. Effective and Efficient Distributed Temporal Graph Learning through Hotspot Memory Sharing is listed as a related work. The paper concludes by presenting DepBloom, a scalable multi-GPU framework for TGNN training. 54. Hongkuan Zhou, Da Zheng, Israt Nisa, Vasileios Ioannidis, Xiang Song, and George Karypis. 2022.

1. The system can achieve a speedup of 2.84× and up to 5.16× over state-of-the-art methods, while maintaining comparable or better model accuracy by implementing DepBloom, which combines temporally bounded spatial partitioning, intra-batch memory refinement, and critical-path optimization.

2. The improved system can handle dynamic graph learning tasks with superior performance by addressing the core challenges of multi-GPU training: temporal dependency forces a trade-off between parallelism and training quality and limiting scalability due to expensive crossGPU memory synchronization.

3. The system can reduce communication overhead by employing temporally bounded spatial partitioning to reduce unnecessary cross-GPU synchronization and critical-path optimization to prioritize urgent memory updates while deferring non-critical work during multi-GPU pipeline execution.

4. The improved system can recover lost short-range temporal information within large batches through intra-batch memory refinement, which is achieved by letting micro-batches execute in parallel and using a module to compute the refined memory as "s(i)v = COMB UPDT s−v, m(i)v, REFINE E(<i)v."

5. The system can enhance pipeline execution efficiency by implementing critical-path optimization for multi-GPU pipeline execution, which ensures that only critical dependencies are synchronized by prioritizing urgent updates and removing non-urgent operations from the critical path.

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.