[论文解读] TGL: A General Framework for Temporal GNN Training on Billion-Scale Graphs

Hongkuan Zhou, Da Zheng+4 more|arXiv (Cornell University)|Mar 28, 2022

Advanced Graph Neural Networks被引用 11

一句话总结

TGL 是一个用于在百亿规模图上进行大规模时序 GNN 训练的统一框架，通过时序采样器、节点内存、邮箱和消息传递引擎，实现了对多种 TGNN 架构的高效训练。它实现了最高达 13 倍的加速，可在 4 块 GPU 上于 1–10 小时内训练包含百亿条边的图，采用如随机分块调度和 Temporal-CSR 等新技术，以实现可扩展性和高性能。

ABSTRACT

Many real world graphs contain time domain information. Temporal Graph Neural Networks capture temporal information as well as structural and contextual information in the generated dynamic node embeddings. Researchers have shown that these embeddings achieve state-of-the-art performance in many different tasks. In this work, we propose TGL, a unified framework for large-scale offline Temporal Graph Neural Network training where users can compose various Temporal Graph Neural Networks with simple configuration files. TGL comprises five main components, a temporal sampler, a mailbox, a node memory module, a memory updater, and a message passing engine. We design a Temporal-CSR data structure and a parallel sampler to efficiently sample temporal neighbors to formtraining mini-batches. We propose a novel random chunk scheduling technique that mitigates the problem of obsolete node memory when training with a large batch size. To address the limitations of current TGNNs only being evaluated on small-scale datasets, we introduce two large-scale real-world datasets with 0.2 and 1.3 billion temporal edges. We evaluate the performance of TGL on four small-scale datasets with a single GPU and the two large datasets with multiple GPUs for both link prediction and node classification tasks. We compare TGL with the open-sourced code of five methods and show that TGL achieves similar or better accuracy with an average of 13x speedup. Our temporal parallel sampler achieves an average of 173x speedup on a multi-core CPU compared with the baselines. On a 4-GPU machine, TGL can train one epoch of more than one billion temporal edges within 1-10 hours. To the best of our knowledge, this is the first work that proposes a general framework for large-scale Temporal Graph Neural Networks training on multiple GPUs.

研究动机与目标

• 解决在大规模动态图上训练时序 GNN 缺乏通用、可扩展框架的问题。
• 通过简单的配置文件，将多种 TGNN 架构统一于单一训练框架之下。
• 克服大规模批量训练中的挑战，如过时的节点内存和百亿规模图上的低效采样。
• 实现对大规模真实世界数据集的 TGNN 评估，超越小型基准测试。
• 提供一个可投入生产的、支持多 GPU 的系统，适用于工业界和学术界使用。

提出的方法

• 设计一种时序 CSR（Temporal-CSR）数据结构，以高效地存储和访问动态图中的时间有序边。
• 开发一种并行时序采样器，通过利用时间顺序和多核处理加速邻居采样。
• 提出一种随机分块调度技术，通过将训练批次打乱为更小、时间上一致的块，减少大规模批量训练中的依赖损失。
• 集成节点内存模块和邮箱，以在时间步之间维护和更新动态节点表示。
• 设计一种消息传递引擎，通过配置文件组合上述组件，支持多种 TGNN 变体。
• 通过 NCCL 和 CUDA 流实现多 GPU 训练，支持数据传输与计算的重叠。

实验结果

研究问题

• RQ1 统一框架能否高效地在百亿规模动态图上训练多种时序 GNN 架构？
• RQ2 当节点内存变得过时时，如何使大规模批量训练实现可扩展性和准确性？
• RQ3 在大规模时序图上，通过优化数据结构和并行采样能获得多大的性能提升？
• RQ4 与小规模基准相比，TGNN 在真实世界百亿边数据集上的表现如何？
• RQ5 多 GPU 训练在多大程度上能实现大规模时序 GNN 训练的可扩展性？

主要发现

• TGL 在链接预测和节点分类任务上，相较于开源基线模型，平均实现了 13 倍的加速，同时保持或提升了准确性。
• 与基线采样器相比，时序并行采样器在多核 CPU 上实现了 173 倍的加速。
• 在 4 块 GPU 的机器上，TGL 可在 1–10 小时内完成超过百亿条时序边的一轮训练，具体时间取决于模型。
• 随机分块调度有效缓解了时间依赖性的损失，使得在大批次大小（如 4800）下也能收敛，而基线方法在此情况下会失败。
• 多 GPU 训练在 JODIE 上实现 2.74 倍加速（4 块 GPU），在 TGN 上实现 2.25 倍加速（4 块 GPU），展现出强大的可扩展性。
• 基于内存的模型（如 TGN）在 GDELT 数据集上单轮训练耗时不足 30 分钟，而非基于内存的模型则需要超过 3 小时。