[论文解读] GraphTheta: A Distributed Graph Neural Network Learning System With Flexible Training Strategy

Yongchao Liu, Houyi Li+9 more|arXiv (Cornell University)|Apr 21, 2021

Advanced Graph Neural Networks参考文献 35被引用 5

一句话总结

GraphTheta 是一个分布式的图神经网络（GNN）学习系统，通过用户自定义的神经网络算子，在基于顶点中心的图处理框架中实现可扩展、灵活的训练，适用于大规模图数据。它引入了 NN-TGAR 抽象和混合并行执行机制，支持多种训练策略，包括集群批量训练，在仅使用基于 CPU 的虚拟机的情况下，对包含 14 亿个节点和 41 亿条边的工业规模数据集，训练速度相比 GraphLearn 最快提升 30.56 倍，相比 DistDGL 提升 2.02 倍。

ABSTRACT

Graph neural networks (GNNs) have been demonstrated as a powerful tool for analyzing non-Euclidean graph data. However, the lack of efficient distributed graph learning systems severely hinders applications of GNNs, especially when graphs are big and GNNs are relatively deep. Herein, we present GraphTheta, the first distributed and scalable graph learning system built upon vertex-centric distributed graph processing with neural network operators implemented as user-defined functions. This system supports multiple training strategies and enables efficient and scalable big-graph learning on distributed (virtual) machines with low memory. To facilitate graph convolutions, GraphTheta puts forward a new graph learning abstraction named NN-TGAR to bridge the gap between graph processing and graph deep learning. A distributed graph engine is proposed to conduct the stochastic gradient descent optimization with a hybrid-parallel execution, and a new cluster-batched training strategy is supported. We evaluate GraphTheta using several datasets with network sizes ranging from small-, modest- to large-scale. Experimental results show that GraphTheta can scale well to 1,024 workers for training an in-house developed GNN on an industry-scale Alipay dataset of 1.4 billion nodes and 4.1 billion attributed edges, with a cluster of CPU virtual machines (dockers) of small memory each (5$\sim$12GB). Moreover, GraphTheta can outperform DistDGL by up to $2.02 imes$, with better scalability, and GraphLearn by up to $30.56 imes$. As for model accuracy, GraphTheta is capable of learning as good GNNs as existing frameworks. To the best of our knowledge, this work presents the largest edge-attributed GNN learning task in the literature.

研究动机与目标

• 解决现有 GNN 框架在训练大规模、深层或高度偏斜图（尤其是高阶度节点）时的可扩展性限制。
• 克服小批量策略中因子图爆炸导致的分布式 GNN 训练中内存和性能瓶颈。
• 在统一、可扩展的系统架构中支持灵活的训练策略，包括集群批量训练。
• 在无需 GPU 或专用硬件的通用 CPU 集群上，实现高效、低内存的 GNN 训练。
• 弥合分布式环境中传统图处理与深度学习工作负载之间的语义与架构鸿沟。

提出的方法

• 在基于顶点中心的分布式图处理引擎中将 GNN 实现为用户自定义函数（UDFs），以实现图计算与深度学习框架的解耦。
• 提出 NN-TGAR（面向递归的神经网络目标图抽象），一种新型抽象，统一图处理与 GNN 操作，简化图卷积的实现。
• 采用混合并行执行：每个训练批次由一组工作进程以分布式方式处理，动态结合数据并行与模型并行。
• 设计一种集群批量训练策略，通过协调地在多个节点上处理子图，减轻内存压力并提升收敛速度。
• 优化随机梯度下降（SGD），利用支持高效邻居采样与跨工作进程梯度聚合的分布式图引擎。
• 采用模块化、可扩展的系统架构，支持无缝集成新的训练策略与图处理原原子。

实验结果

研究问题

• RQ1 基于顶点中心图处理模型构建的分布式 GNN 训练系统，是否能在每个节点内存有限的情况下，实现对工业规模图的高可扩展性与高性能？
• RQ2 所提出的 NN-TGAR 抽象相比传统框架，在实现图卷积时是否显著提升了表达能力与效率？
• RQ3 在大规模、带边属性的图上，GraphTheta 在不同训练策略（尤其是集群批量）下的性能与可扩展性如何？
• RQ4 在真实世界数据集上，GraphTheta 是否能在训练速度与模型泛化能力方面超越现有分布式 GNN 框架（如 DistDGL 和 GraphLearn）？
• RQ5 纯 CPU 系统在无 GPU 加速的情况下，能否在大规模 GNN 训练中实现具有竞争力的性能？

主要发现

• GraphTheta 在基于 CPU 的虚拟机集群上可有效扩展至 1,024 个工作节点，展现出大规模 GNN 训练的强横向可扩展性。
• 在蚂蚁金服数据集（14 亿个节点，41 亿条边）上，训练四层 GCN 时，GraphTheta 相比 GraphLearn 最快提升 30.56 倍，相比 DistDGL 提升 2.02 倍。
• 集群批量训练策略在蚂蚁金服数据集上实现了最快收敛与最佳泛化性能，优于标准小批量与全局批量方法。
• 在 Reddit 等多个基准数据集上，GraphTheta 的模型准确率与现有框架（如 DistDGL 和 GraphLearn）相当或更优。
• 系统在高度偏斜的图上表现出色，即使节点邻居数达数十万，也能通过高效的批量管理避免子图爆炸。
• 尽管仅运行在 CPU 上，GraphTheta 仍实现了具有竞争力的训练吞吐量，表明其在公共云基础设施上具备低成本、生产级部署的强大潜力。