蚂蚁金服人工智能部论文《AGL:可扩展工业图机器学习系统》，处理十亿节点千亿边图数据的GNNs训练推理

2020 年 3 月 9 日 专知

【导读】图机器学习Graph Machine Learning 是当下的研究热点，受到学术界和工业界的关注。最近，蚂蚁金服人工智能部在arXiv上公布了最新研发的图机器学习系统-AGL论文-《AGL: a Scalable System for Industrial-purpose Graph Machine Learning》，该系统具有完整的GNNs训练和推理功能，能够处理十亿节点和千亿边的图数据处理，值得关注。

地址：

https://arxiv.org/abs/2003.02454

基于图的机器学习已经成为强大的图形数据学习工具。然而，由于图数据固有的依赖性，工业社区很难利用图神经网络(GNNs)等技术来大规模地解决实际问题。因此，我们不能简单地用经典的学习系统来训练GNN，例如假设数据并行的参数服务器。现有系统将图形数据存储在内存中，以便在单机或远程图形存储中进行快速访问。主要的缺点有三方面。首先，由于内存容量或图数据存储与workers之间的带宽的限制，它们无法进行扩展。其次，它们需要在没有充分利用MapReduce等成熟基础设施(这些基础设施保证了良好的系统属性)的情况下，对图形存储进行额外的开发。三是注重训练，忽略了对图的推理优化，使其成为一个非集成系统。

本文设计了一个可扩展的容错集成系统AGL，该系统具有完整的GNNs训练和推理功能。我们的系统设计遵循基于GNNs计算的消息传递方案。我们设计生成k-hop邻居，即每个节点的信息完备子图，并通过简单地合并来自边缘内邻居的值和通过MapReduce将值传播到边缘外邻居来进行推理。此外，k-hop邻居包含每个节点的信息完备子图，因此由于数据独立性，我们只需在参数服务器上进行训练。我们的AGL系统是在成熟的基础设施上实现的，可以在14小时内完成一个2层图注意力网络在一个数十亿节点、千亿边的图上的训练，并在1.2小时内完成推理。

介绍

近年来，机器学习在图结构数据研究中的应用越来越受到工业界和学术界的重视。图机器学习(缩写为GML)不仅在传统的图挖掘的任务取得成功(例如,节点分类(b1, b2, b3, b4),链接属性预测(b5)和图属性预测[b6, b7]), 也带来了巨大的改善,其他领域的任务(例如,知识图谱(b8、b9), NLP (b10), 计算机视觉(b11, b12),等等)。此外，越来越多的互联网公司将GML技术应用于解决各种行业问题，并取得了巨大的成功(如推荐[b14, b14_2]，营销[liu2019graph]，欺诈检测[b15]，贷款违约预测[b16]等)。

为了使用图机器学习技术处理工业规模图数据来解决实际问题，我们需要构建一个具有可扩展性、容错性和全功能训练/推理工作负载完整性的学习系统。然而，由于数据依赖性，图机器学习任务的计算图与传统学习任务有着本质的区别。即在现有经典参数服务器框架[b22]中，假设数据是并行的，每个样本的计算图独立于其他样本，而图学习任务中每个节点的计算图依赖于k-hop 节点的邻居。图数据学习任务中的数据依赖性使得我们无法再将训练样本或推理样本存储在磁盘中，并通过管道进行访问[b22]。相反，我们必须将图形数据存储在内存中，以实现快速的数据访问。这使得我们不能简单地基于现有的参数服务器框架为图数据学习任务构建一个学习和推理系统，而现有的参数服务器框架只是简单地在参数服务器中维护模型一致性，并在每个工作人员中并行处理数据。

一些公司努力为各种GML技术设计巧妙的系统架构。Facebook提出了PyTorch-BigGraph (PBG) [b17]，这是一个大规模的网络嵌入系统，旨在从多关系数据中产生无监督节点嵌入。然而，PBG并不适合大量的真实场景，在这些场景中，图具有丰富的节点和边缘属性(称为属性图)。最近，利用深度学习技术，图神经网络(GNN)能够生成属性图的高质量表示，或者在带标记节点或边的属性图上构建端到端的监督模型，从而成为更通用的GML技术。例如，深度图库(Deep Graph Library, DGL) [b18]、PyTorch Geometric (PyG) [b24]和AliGraph [b19]已经被开发出来用于训练大规模属性图上的图神经网络。其中，DGL和PyG被设计成一个单机系统，以处理内存中的工业规模的图形，该系统基于一个巨大的机器，例如拥有2TB内存的AWS x1.32xlarge。作为一个分布式系统，AliGraph实现了分布式内存中的图形存储引擎，这需要在训练GNN模型之前进行独立部署。

表1:不同GML系统的图规模

然而，真实的工业图数据可能是巨大的。Facebook的社交图包含超过20亿个节点和超过1万亿个边[b17, b20]。蚂蚁金服的异构金融图包含数十亿个节点和数千亿条具有丰富属性信息的边，阿里巴巴的电商图也一样。表1总结了由几个当下的GML系统报告的图规模。通过计算与这些节点和边相关的特征，这种规模的图形数据可能产生100 TB的数据。这些数据不能存储在像DGL这样的单一机器中。此外，图数据存储引擎存储与节点和边缘相关的图数据和特性，Worker之间的通信可能非常庞大。例如，假设一个具有1000个节点和10000条边的子图的批处理，这可能会导致在图存储和worker之间通信1MB块，如果我们不能在管道中访问数据，这是无法容忍的。此外，这需要一个结构良好、带宽足够大的网络。

总而言之，首先，现有的学习系统的工业设计需要图形数据的内存存储，要么是在无法处理实际工业规模的图形数据的单个巨型机器中，要么是在定制的图形存储中，这会导致图形存储与workers之间的大量通信。这使得它们不能扩展到更大的图形数据。其次，它们不能很好地利用传统的基础设施(如MapReduce或参数服务器)来实现容错。第三，现有的框架大多注重图学习模型的训练，而忽略了系统的完整性，例如在部署图机学习模型时，推理任务的优化。

考虑到这些问题，我们构建了AGL (Ant Graph machine Learning system)，这是一个用于工业图学习的集成系统。我们的系统设计的核心思想是基于基于图神经网络计算图的消息传递(合并和传播)方案。

在训练图神经网络的阶段，我们提出了构造k-hop邻域的方法，该方法通过将邻域内的边合并，并将合并后的信息沿外边传播给邻域，从而根据消息传递为计算每个节点的k-hop嵌入提供信息完备的子图。将原始图分解成极小的子图，即k-hop邻域，其优点是每个节点的计算图又独立于其他节点。这意味着我们仍然可以享受来自经典参数服务器框架的容错特性和灵活的模型一致性，而无需额外维护图数据存储[b19]。

在图神经网络的推理阶段，我们提出将一个训练良好的K层图神经网络分割成K个切片和一个与预测模型相关的切片。对于切片，我们首先合并每个节点的内边缘邻居的第k层嵌入，然后将嵌入传播到它们的外边缘邻居，k从1到k。

我们在训练和推理中抽象出所有的消息传递方案，并简单地使用MapReduce [dean2008mapreduce]实现它们。由于MapReduce和参数服务器都是在工业企业中作为基础设施开发的，所以我们的图机器学习任务系统仍然可以受益于容错和可靠性等特性，即使是使用廉价且广泛使用的普通机器。此外，与基于DGL和AliGraph等架构的推理相比，我们的推理实现最大限度地利用了每个节点的嵌入，从而显著提高了推理作业。此外，我们还提出了一些技术来加速浮点运算在训练过程中从模型级到操作级的运算。因此，与DGL/PyG相比，我们成功地在单机中加速了GNNs的训练，并在实际产品场景中通过CPU集群实现了近乎线性的加速。

值得注意的是，AGL在处理6.23×10^9个节点、3.38×10^11条边的图时，可以在14小时内完成1.2×108个目标节点的2层GAT模型的训练(7个epochs直到收敛，100个workers)，仅用1.2小时完成对整个图的推理。据我们所知，这是迄今为止最大的图嵌入应用，并证明了我们的系统在实际工业场景中的高可伸缩性和效率。

AGL架构

图1: AGL的系统架构

AGL的设计原则是基于基于GNNs计算的消息传递方案。也就是说，我们首先合并来自每个节点的内边缘邻居的所有信息，然后通过外边缘将这些合并后的信息传播到目标节点。我们反复将这一原理应用到训练和推理过程中，开发了GraphFlat和GraphInfer。基本上，GraphFlat是在训练过程中生成独立的K-hop邻域，而GraphInfer是在训练良好的GNN模型下推断节点的嵌入。

基于这种动机和设计原则，提出的AGL利用几个强大的并行架构(如MapReduce和Parameter Server)，用精心设计的分布式实现构建每个组件。因此，即使部署在具有相对较低计算能力和有限内存的集群上，AGL也可以获得与一些最先进的系统相当的效率和更高的效率。此外，它还具有在具有数十亿节点和千亿边的工业规模图上执行全功能图机器学习的能力。