研究摘要
随着图神经网络（GNN）向动态图数据扩展，时序图神经网络（TGNN）在处理动态图数据方面展现了卓越的能力。然而，在分布式 TGNN 训练中，高效处理时序依赖性带来的大量跨设备通信成为关键挑战，这通常导致显著的冗余数据传输和高昂的通信开销。现有系统难以有效消除数据重用与传输中的冗余问题，因此在分布式环境下表现出严重的通信瓶颈。为此，本研究提出了 Sven，一种联合设计的算法与系统库，专为多 GPU 平台上的 TGNN 训练加速而构建。Sven 利用 TGNN 模型的依赖模式，开发了无冗余的图组织方式，从根本上减少了冗余数据传输。同时，我们针对设备间通信不平衡问题，将图划分问题形式化为最小化最大通信代价问题，并证明其为 NP 难问题。为此，我们提出了一种近似算法 Re-FlexBiCut 来解决这一问题。此外，Sven 结合预取技术、自适应微批次流水线以及异步流水线机制，构建了分层流水线方法以缓解通信开销。
研究背景
图1： 时序图神经网络训练流程
图作为一种强大的数据结构广泛应用于多个领域，能够高效地将真实世界的对象建模为关系数据结构。近年来，图神经网络（Graph Neural Networks, GNNs）因其在处理图数据方面的卓越性能受到广泛关注，广泛应用于节点分类、链路预测和图分类等任务中。尽管诸如 GCN 和 GraphSAGE 等方法在静态图（具有固定节点和边的图）上取得了显著进展，但它们无法处理现实世界中不断变化的动态图数据。为了克服这一限制，时序图神经网络（Temporal Graph Neural Networks, TGNNs）应运而生。TGNN 不仅能够捕捉动态图的时间信息，还可以同时学习其拓扑关系。研究表明，TGNN 在性能上显著优于传统的静态 GNN 方法，为动态图分析任务提供了更强大的工具支持。
然而，动态图中丰富的信息处理通常需要大规模并行和分布式计算来支持 TGNN 的高效训练。为捕获时序依赖性，许多 TGNN（如 TGN、APAN 和 JODIE 等）采用基于记忆的架构，通过节点记忆和消息模块总结每个顶点的历史行为。这种方法在扩展 TGNN 训练时引发了两个关键的性能问题。第一，TGNN 需要根据所有顶点的交互更新最新的时序依赖性，导致大量的冗余依赖数据；第二，在大型 GPU 集群中，时序依赖性和模型参数分布在多个设备上，需要通过全互通信（如 all-to-all 和 all-gather）来分发和聚合依赖数据，同时通过 all-reduce 同步模型参数的梯度。这些通信操作显著增加了训练时间开销，据测量，通信时间在 TGNN 训练中可占据 70% 以上。因此，如何通过消除冗余数据和利用层次化流水线技术降低通信开销成为研究的核心动机。
研究现状
目前，已有多种 GNN 系统框架支持高效灵活的静态图神经网络训练，例如 PyG 和 AGL。然而，这些框架在动态图支持上较为有限甚至完全缺失。为了提高动态图训练的效率，一些专门设计的框架被提出。例如，DGL 提供了一些基础接口以支持动态图训练，但这些接口效率低下且缺乏对分布式环境的兼容性。TGL 是一个支持大规模离线 TGNN 训练的框架，能够在单节点多 GPU 环境下高效训练多种 TGNN 模型。然而，TGL 由于频繁的主机与设备之间内存传输，面临显著的通信开销问题。ESDG 提出了一种基于图差分的算法，能够减少通信流量并扩展 TGNN 训练至多节点多 GPU 系统。然而，ESDG 仅支持基于离散时间动态图（DTDG）的 TGNN 训练，无法推广到更通用的连续时间动态图（CTDG）模型训练。
最近，依赖缓存（DepCache）和依赖通信（DepComm）机制被提出，用于分布式 GNN 训练中的依赖性维护。然而，我们的实验结果表明，这两种机制在 TGNN 训练中效率较低，其主要瓶颈在于大量的跨设备通信开销。此外，由于 TGNN 的依赖性管理对通信流量和计算性能的要求更高，目前的方案仍无法满足分布式环境下的高效训练需求。
综上所述，现有框架和方法在支持动态图尤其是连续时间动态图的训练上仍存在显著局限性，亟需提出新的方法以优化 TGNN 训练的冗余数据处理和通信开销，从而实现更高效的分布式 TGNN 训练方案。
研究方法
图2： Sven的系统架构：算法-系统联合设计
本研究提出了 Sven，一个为多节点多 GPU 系统上高性能分布式时序图神经网络（TGNN）训练设计的算法与系统联合框架，专门针对连续时间动态图（CTDGs）的 TGNN 训练需求。Sven 从整体视角和协同方法出发，系统性地解决了 TGNN 训练中冗余依赖性和通信与计算时间消耗失衡的问题。首先，我们分析了 TGNN 模型在处理时序依赖性时的行为，发现模型需要在当前状态下分发依赖性用于计算，并在计算后聚合最新依赖性。由于一个顶点可能多次与其他顶点发生交互，这种操作在分发和聚合阶段会产生大量冗余数据。我们对多个动态图数据集的冗余比例进行了测量，结果表明通过从原始图中提取无冗余图，能够有效减少时序依赖性通信中的数据冗余。此外，我们改进了现有的 DepCache 和 DepComm 机制以维护分布式 TGNN 训练中的时序依赖性通信。针对时序依赖性通信（包括 DepCache 的 all-gather 操作和 DepComm 的 all-to-all 操作）中产生的大量数据冗余，我们提出了从源头减少通信量的优化方案。为进一步降低通信开销的影响，我们提出了一种自适应微批次流水线方法，以减少 all-to-all 操作引入的系统开销，同时设计了异步流水线方法，通过隐藏 all-reduce 通信来提升整体性能。此外，我们的研究表明，现有的图划分方法会导致设备间通信负载不平衡，从而显著影响分布式 TGNN 训练的性能。为解决通信负载不平衡问题，我们构建了更严谨的问题形式化，并设计了一个针对通信负载优化的近似算法 Re-FlexBiCut。该方法通过为每个分区构建源/汇图，并允许顶点在分区间重新分配，从而实现接近平衡的最小切割方案，有效提升了分区间的通信平衡性，最终显著优化了分布式 TGNN 训练的性能。
研究成果
图3： 端到端训练时间对比
本研究提出的 Sven 框架在分布式时序图神经网络（TGNN）训练领域取得了显著成果，成功解决了冗余依赖性数据和通信负载不平衡问题，为大规模动态图数据的高性能训练提供了全新方案。作为首个面向内存驱动 TGNN 训练的全面优化方案，Sven 在 64 GPU 集群上的实验结果表明，其训练速度相较于现有最先进方法提升了 1.9 至 3.5 倍，同时通信效率最高提升 5.26 倍，并减少了高达 59.2% 的通信不平衡问题，为大规模 TGNN 训练提供了一种高效的分布式解决方案。
相关研究成果获得了高性能计算领域的广泛认可，并发表在以下顶级会议和期刊中：
1. 论文《Redundancy-Free High-Performance Dynamic GNN Training with Hierarchical Pipeline Parallelism》 发表在高性能计算顶会 HPDC 23 (ACM International Symposium on High-Performance Parallel and Distributed Computing)（CCF-B），并荣获 最佳论文亚军。
2. 论文《Redundancy-free and load-balanced TGNN training with hierarchical pipeline parallelism》 发表在 CCF-A 类顶级期刊 IEEE TPDS（Transactions on Parallel and Distributed Systems）上。