Amazon Launches DistTGL: A Memory-Based Temporal Graph Neural Network Framework

Recently, researchers have developed Temporal Graph Neural Networks (TGNNs) that leverage temporal information in dynamic graphs, building upon the success of Graph Neural Networks (GNNs) in learning static graph representations. TGNNs have demonstrated superior accuracy in various downstream tasks on dynamic graphs, such as temporal link prediction and dynamic node classification, significantly outperforming static GNNs and other traditional methods. In dynamic graphs, each node has more relevant events over time. When this number is high, TGNNs cannot fully capture the history using time-based attention aggregation or historical neighbor sampling techniques. Researchers have created Memory-based Temporal Graph Neural Networks (M-TGNNs), which store node-level memory vectors to summarize independent node histories and compensate for the lost history. Although M-TGNNs have achieved success, their poor scalability makes them difficult to implement in large-scale production systems. Due to the time dependency generated by auxiliary node memory, training mini-batches must be short and arranged in chronological order. Utilizing data parallelism in M-TGNN training is particularly challenging and involves two aspects: 1. Increasing batch size alone leads to information loss and the loss of temporal dependency information between events. 2. All trainers must access and maintain a unified version of node memory, resulting in a significant amount of remote traffic in distributed systems. A new study by the University of Southern California and AWS provides DistTGL, a scalable and efficient method for training M-TGNNs on distributed GPU clusters. DistTGL improves the current M-TGNN training system in three aspects: 1. Model: By introducing more static node memory, the accuracy and convergence speed of M-TGNNs' node memory are improved. 2. Algorithm: To address accuracy loss and communication overhead in distributed settings, the team provides a novel training algorithm. 3. System: To reduce the overhead related to mini-batch generation, they have developed a system optimized using prefetching and pipelining techniques. DistTGL significantly outperforms previous methods in terms of convergence speed and training throughput. DistTGL is the first work to extend M-TGNN training to distributed GPU clusters. DistTGL is publicly available on GitHub. Based on the unique characteristics of M-TGNN training, the researchers propose two innovative parallel training methods - epoch parallelism and memory parallelism, which enable M-TGNNs to capture the same number of dependent graph events on multiple GPUs as on a single GPU. They provide heuristic recommendations for selecting the optimal training settings based on the dataset and hardware features. The researchers serialize memory operations on node memory and efficiently execute them through a separate daemon process, eliminating complex and expensive synchronization to overlap mini-batch creation and GPU training. In experiments, when scaled to multiple GPUs, DistTGL improves convergence speed by more than 10 times compared to state-of-the-art single-machine methods.