GNN

前言 工业界的图规模都非常大，少说也是上千万的顶点+上亿的边，单机训练不现实，必须借助多机分布式训练。然而目前主流的图训练框架PyG、DGL对图的多机分布式训练支持都不太好。工业界好像阿里的Euler、百度的PGL可以支持分布式训练。今天介绍一下亚马逊DGL针对分布式训练所做的优化。 摘要 GNN广泛应用在推荐、搜索、风控等领域，在这些领域，图的规模往往非常大，有数以亿计的顶点和万亿的边。为支持大规模图的分布式训练，本文提出了DistDGL，它能以mini-batch的方式在多机上进行分布式训练。DistDGL基于DGL框架，它将图数据分布在多台机器上，并基于数据分布，将计算也分布在多台机器上（owner-compute rule）。DistDGL以同步更新的方式进行训练。为了减小分布式训练的通信开销，DistDGL使用一个高效、轻量的图分割算法对图进行分割，在分割时设计了多个负载均衡约束，使得每个分割的子图达到较好的负载均衡。此外，为了减小跨机器的通信，DistDGL在每个子图中保留了halo nodes（正文会介绍到），并且使用了稀疏embedding更新策略。这些优化策略使得DistDGL在分布式训练时能达到较好的高并行效率和内存可扩展性。实验结果表明，在分布式训练时，随着计算资源的增大，DistDGL的训练速度可以线性增长。在16台机器组成的分布式环境中，DistDGL仅用13秒就可以完成1亿节点+30亿边的一个epoch的训练。DistDGL是DGL的一部分，已开源在：https://github.com/dmlc/dgl/tree/master/python/dgl/distributed。 简介 GNN很有用，但是现实世界中的网络都很大，比如Facebook的社交网络、Amazon的用户商品关系网络等。 GNN分布式训练的难点： GNN中每个训练样本（顶点）不是独立的，是相互依赖的，比如为了训练顶点A，必须采样A的邻居，随着GNN层数的增大，采样的邻居数目呈指数上升。而CV、NLP中每条样本是相互独立的。 GNN的多机分布式训练的通信数据主要是图数据（顶点和边，及其属性），而CV、NLP的分布式训练的通信数据主要是网络参数、梯度等，需要通信的数据类型不同，导致CV、NLP的分布式训练优化技术无法直接迁移到GNN的分布式训练中。 此外，神经网络大多采用同步更新的分布式训练策略（难道不是异步？），因此需要尽量做到不同机器或者worker的负载均衡。由于图的特殊结构，不同顶点的度差异很大，即不同子图的负载差异很大，所以如何实现GNN训练时的负载均衡，也是一个难点。 背景介绍 GNN 以消息传递的方式来解读GNN，每一层GNN可以用下面的公式来概括。\(\mathbf{h}_v^{(l+1)}\)和\(\mathbf{h}_v^{l}\)分别表示节点\(v\)在第\(l+1\)层和第\(l\)层的向量表示。\(f\)表示节点\(v\)和每个邻居\(u\)计算消息；\(\oplus\)表示邻居聚合函数；\(g\)用来更新节点表示。 作者将GNN的参数分为两部分，一部分是网络参数，即上面的\(f\)、\(\oplus\)和\(g\)。另一部分是节点本身的embedding参数，对于transductive模型来说，节点本身有embedding，故有这部分参数；但对于inductive模型，节点本身没有embedding参数，节点的embedding表示是通过网络参数生成的。 为了区分这两部分参数，作者将网络参数称为稠密参数dense parameters，所有dense参数在每个mini-batch都需要被全部更新。作者将顶点本身的embedding参数称为稀疏参数sparse parameters，每个mini-batch只需要更新该batch涉及到的顶点的稀疏参数即可。 Mini-batch training GNN进行mini-batch训练时的基本流程如下： 从训练集中随机采样N个顶点，这部分顶点称为target vertices 对每个target vertex随机采样最多K个邻居顶点 对每个target vertex，通过聚合其邻居的信息得到target vertex的表示 上述流程是一层GNN的训练过程，如果GNN有多层，则邻居采样的过程会递归进行下去。 方法 DistDGL分布式训练框架 DistDGL的核心可以用上面的图来表示。DistDGL包含三个组件： Trainer，即图中的GNN Training Component，其中放大的GNN Training Component只是右边3个之一的放大图而已。Trainer主要是用来训练的，即进行前向传播和反向传播的。 Sampler，即图中的Sampling Component，用来采样邻居的。 KVStore，即图中的KVStore Component，用来存储顶点和边的特征，以及相应的embedding。 对照背景介绍中的mini-batch training过程，DistDGL的训练过程如下： Trainer随机采样N个顶点作为target vertices Trainer向Sampler请求采样target vertices的邻居 Trainer向KVStore请求target vertices及其邻居的属性信息 Trainer开始分布式训练，并使用AllReduce方式同步更新dense参数（网络参数）；并将sparse参数（embedding）存储到KVStore中 主要优化点 图分割及负载均衡 这是DistDGL最核心的优化点。为了实现高效的分布式训练，DistDGL首先使用METIS图分割算法把图分割成多个子图，不同子图分布式存储在不同机器上；然后把不同子图的计算也分配到存储数据的机器上。做到数据在哪里，计算就在哪里（owner-compute rule），最大程度利用数据和计算的局部性，减小网络通信。 如下图所示，METIS算法以最小割的方式分割图网络，即如果每条边都有不同的权重的话，METIS希望分割的时候切割的边的权重之和最小，由此可以尽量把有密切连接的节点分割到同一个子图中。 我没仔细研究METIS算法，我理解METIS还需要一个约束，即需要分割成多少个子图，或者每个子图最多有多少个顶点之类的。要不然什么都不分割，直接输出全图，则割最小是0。 如果某一条边被分割了，其所连的两个顶点被分割到两个不同的子图中了，称这样的顶点为HALO vertices（通俗理解就是边缘点）。如果需要采样HALO顶点的邻居，则需要跨子图进行采样，涉及到网络通信。为了避免网络通信成为瓶颈，DistDGL会在两个子图中都保留HALO顶点的另一端顶点。在这种情况下，如果只涉及到HALO顶点的一跳采样的话，不需要跨子图通信。DistDGL通过冗余存储HALO顶点，以减小网络通信。由于GNN网络的邻居采样一般只会有2-3跳，所以这种策略应该能避免大部分跨子图通信。 由于同一批数据只需要在开始训练时做一次分割，相对于漫长的N个epoch训练时间来说，分割图的时间开销被分摊了，可以忽略不计。 分割完图之后，DistDGL把不同子图分配到不同机器上。在训练的时候，由于trainer、sampler和KVStore需要互相交换数据，为了提高数据交换效率，DistDGL把属于同一个子图的trainer、sampler和KVStore分配到了同一台机器上，则三者之间的通信可以直接通过共享内存的方式进行内存拷贝，大幅减小了网络通信带来的延时。 如下图所示，同一台机器上的Trainer、Sampler和KVStore是共享内存的。 文中还提到在METIS分割图的时候，增加了很多约束条件，以达到负载均衡的目的。我理解默认METIS在进行分割的时候，可能只保证不同子图的顶点数大致相同，在这个约束下去最小化割。然而子图的顶点数相同，并不代表子图的负载也相同，还涉及到子图中边的数目，不同类型顶点的数目分布等等。因此DistDGL在METIS子图分割时还增加了很多约束条件，使得分割的子图在训练的时候尽量达到负载均衡。 分布式KVStore DistDGL把顶点和边的属性特征及embedding存储在分布式KVStore中，DistDGL开发了自己的分布式KVStore，而不是使用现成的比如Reddis，原因是更方便自定义功能，比如把属于同一个子图的顶点、边、特征存储到同一个机器上，优化了网络传输，实现稀疏embedding的异步更新等。 分布式Sampler Trainer训练和Sampler采样是并行进行的，简单理解就是，Trainer在训练当前Epoch数据的时候，Sampler就已经在异步采样下一个Epoch的数据了，充分利用计算资源，实现流水线作业。类似的，局部采样和远程网络RPC通信也可以overlap“同步”进行，使得局部采样感受不到远程通信的等待时间。 ...

前言 最近的工作涉及到图神经网络，打算系统学习下这方面的内容。首先搜集了相关的教材，发现市面上的教材大多数是罗列论文的形式，不太适合初学者入门。后来找到了斯坦福CS224W这门公开课，打算入坑，一是之前学习过斯坦福CS224N，感觉不错；二是CS224W这门课的老师是GraphSAGE的作者Jure Leskovec，有大佬背书错不了。 CS224W主页：http://web.stanford.edu/class/cs224w/ Winter 2021版主页：http://snap.stanford.edu/class/cs224w-2020/ Winter 2021版视频：https://www.youtube.com/playlist?list=PLoROMvodv4rPLKxIpqhjhPgdQy7imNkDn，Jure Leskovec是斯洛文尼亚人，英语不是很标准，建议打开YouTube的字幕。 背景介绍 图（Graph）是描述实体（entity）和关系（relation）的一种通用语言形式，它由节点（vertex或node）和连接节点的边组成，很多数据类型都可以用图的形式来描述。 图1 图及其应用实例 目前常见的图有两类： 第一类是网络（network），也称为自然图，例如： 社交网络，全球70亿人形成一个大网络 通信网络，例如通过电话、邮件、交易等形成的网络 生物医药网络，例如基因、蛋白质之间形成的网络 大脑中的成千上万的神经元形成的网络 第二类是通过抽象表示形成的图，例如 人工组织形成的信息网络、知识网络 软件中的代码调用形成的网络 分子网络、场景图、基于粒子的物理模拟等 现有的机器学习工具箱主要针对图像、文本和语音，对图的机器学习处理工具相对较少，因为图是不规则的数据，难以处理。对图的处理主要有以下难点： 图不是欧几里得数据结构，没有固定的大小和拓扑结构 图上的节点没有固定的顺序，也没有参考点，是去中心化的 图会随着时间动态变化，并且图中常常会融合多模态信息 本课程的两个重点： Deep learning in graphs，即图上的深度学习算法 Representation learning，即图表示学习，将图中的节点嵌入到一个低维稠密向量中，使得网络中相似节点的embedding距离接近 本课程的主要内容包括： 传统方法：Graphlets，Graph Kernels 节点嵌入方法：DeepWalk，Node2Vec 图神经网络：GCN，GraphSAGE，GAT，Theory of GNNs 知识图谱：TransE，BetaE 图上的深度生成网络 图在生物医药，科学和工业上的应用 图机器学习应用 图可以有很多应用场景，这些应用可以分为节点水平的（nodel level）、边水平的（edge level）、子图水平的（subgraph level）和图水平的（graph level）。下面逐一举例： Node-level：节点分类（node classification），例如预测节点的属性。节点回归？例如AlphaFolde使用GNN预测每个氨基酸在三维空间中的位置坐标，从而预测蛋白质的结构。感觉和GNN关系不太大吧？具体得看论文了。 Edge-level：链接预测（link prediction），预测两个节点之间是否存在边。例如在推荐系统中，预测user是否会购买item等。另外还可以用于预测药物的副作用，例如任意两种药组合吃，是否会产生副作用，产生哪种副作用，都是针对边的任务。 Sub-graph level：地图导航，预测预期到达时间（ETA）。DeepMind和Google Maps合作的一个工作，很有意思：https://www.deepmind.com/blog/traffic-prediction-with-advanced-graph-neural-networks。简单来说，把每条路分段（supersegment），每段表示成一个点，一条路的相邻段（点）连边，交叉路口的段（点）连边。通过GNN的消息传递，一条路的拥堵信息，可以传递到相邻的路。很自然的想法，也符合实际情况，比如在这条路拥堵了，司机可能就会走相邻的路，进而会影响相邻的路的ETA。问题是，GNN对图很敏感，不同地区、地段的路网图差异很大，有的路网小，有的路网大，因此不同training run之间的方差很大。一开始想到用lr decay来缓解。后来使用MetaGradients让模型自动调整学习率。使用多个loss，多目标学习防止过拟合。 Graph-level：例如新药发现：节点是原子、边是各种键，生成一个graph，就是一种新的复合物。物理模拟：动态图，节点表示粒子，有属性比如速度、动量，然后下一个时刻有新的位置，不断进化变化，类似RNN，可以模拟出粒子的动态变化过程。 图2 图机器学习应用场景 图的表示方法 构成图的基本要素包括顶点集合N和边集合E，可以用\(G(N,E)\)来表示一张图。 根据边是否有方向，可以将图分为无向图和有向图，无向图即图中的边没有方向，有向图即图中的边有方向。 对于无向图G，每个顶点的度就是该顶点所连边的数目，由于一条边连接了两个顶点，贡献了2个度，所以所有顶点的平均度数=2E/N。 对于有向图，顶点的度可分为入度和出度，如图3所示，顶点C的入度为2，出度为1。所有顶点的平均入度=平均出度=E/N。如果某个顶点的入度为0，则称该顶点为源点，例如顶点G；如果某个顶点的出度为0，则称该顶点为槽点（sink），就像水槽一样，只进不出；如果某个顶点的入度和出度都为0，则称该顶点为孤立点。 图3 图的表示方法和顶点的度 ...