转载公众号 | DataFunSummit
本文主要介绍飞桨图学习的大模型训练框架。
介绍会围绕下面三点展开：
1. 图学习框架发展趋势
2. 图学习框架的 GPU 加速
3. 工业级大规模图表示学习 PGLBox
分享嘉宾｜黄正杰 百度 资深研发工程师
编辑整理｜王超 智联招聘
出品社区｜DataFun
01
图学习框架发展趋势
首先和大家分享下联邦学习的背景。
1. 什么是图？以及图的来源
首先图是描述复杂世界的一个通用语言。包括我们在互联网里面的一些像社交网络，它是人和人之间的连接。以及近期一些比较火热的像生物计算的一些领域，它是一些化学分子的图的结构。以及像知识图谱，它是一些实体与关系的连接。推荐系统是用户与物品之间的一些关联。
2. 图神经网络的发展趋势
图神经网络这个概念也是在近年来比较火的。之前在建模图的时候，通常会基于一个谱域的概念，就是一个 Spectral-based 的方法。会首先把一些图通过傅立叶变换转化到频域，再去频域上操作。当时可能是因为这个理念比较复杂，并且它对于大图的扩展性比较差，所以关注这方面工作的人比较少。
大概在 2017 年左右，像图卷积网络，还有一些像 Message Passing 的网络出现之后，大家会转向用更简单的架构去看待图网络，例如 Spatial-based 的方法，就是在空域上去理解图神经网络。我们可以类比到二维的图卷积里面去，如果你要建模一个点的元素，它是跟你周围的像素有关的。在图神经网络里面的话就是如果我要建模一个点、它的一个属性以及表示，通常只会跟它的邻居有关，然后就通过一些邻居的计算，我们就可以得到我们中心节点的一个表示。
3. 基于空域（Spatial）的图神经网络
在基于空域的图神经网络里面，我们通常是关注两个点，一个是我们怎么做邻居的聚合去得到中心点的表示；另一个是我们有一张大图了，怎么去通过节点的汇聚得到整个图结构的表达。
4. 基于消息传递的图神经网络
目前主流的图神经网络基本都可以描述成这种叫做消息传递的形式，即 Message Passing。像 DGL 还有 PYG 这些目前比较热门的图神经网络框架，包括我们的 PGL 也是沿用这样基于消息传递的范式去定义图神经网络。
这里的 Message Passing 就是说我们只需要建模邻居是怎么把的消息从源节点发送到一个目标节点，以及目标节点获取到源节点的消息之后怎么做，例如做一个加权的求和或者是做一个均值。通过这样的一个框架，我们目前基本上能实现大部分的图神经网络，也甚至包括现在的 Transformer，假设把 Transformer 理解成一个全连接的图神经网络，其实 Transformer 的结构也是可以通过这种消息传递形式去实现。
5. 飞桨图学习框架 PGL
上图是我们的一个图学习的框架 PGL。大概是在两年前发布的，主要是基于飞桨的 PaddlePaddle 深度学习框架，在上面搭建了一些包括图的引擎，以及图网络的一些编程的接口，上面我们也放了很多预置的模型，供大家去使用。感兴趣的同学可以去我们的Github，可以 pip install 装我们的框架，以及跑一些示例。
6. 大规模图神经网络的训练标准流程
图神经网络大规模的训练，目前是有一套比较成熟的标准流程。特别是大图的情况下，我们通常会沿用 2017 年的一篇论文叫做 GraphSAGE，它主要的理念就是训练在大图里面的某些带标签的节点，我们通常会对某个点的邻居进行采样 sample，后面用 aggregation 就是不停地聚合或者通过消息传递的图神经网络去代表我们中心的一个点。
对于不同的框架，我们都会主要解决三个问题，第一个是我们的图是怎么存的，第二个是我们的子图是怎么采样，包括一些特征是怎么从图里面拉取出来。最后是高层图网络的一些消息传递算法是怎么定义的。
相比于一些大模型，在图网络里面大部分的消耗反而是在图的存储以及特征的拉取上面。比起刚刚说到的一些大模型参数量上，目前模型在纯网络的模块是会比较小的，但即使在这种情况下，我们的图神经网络训练的 GPU 利用率还是很低的。最大的瓶颈就在于 CPU 上图的采样的瓶颈，以及我们从 CPU 到 GPU 的数据传输的瓶颈。
因为我们的图很大，所以通常情况下 GPU 是存不下来的。我们会把图和一些采样的算法实现在内存，最后采样完之后，我们再把它挪到我们的 GPU 去做计算。
7. 硬件的变化：显存越来越大
随着硬件的逐渐地发展，我们也会发现一个趋势，就是现在的显存越来越大，包括是 8 块的 A100 的 80G 我们能在单台机器上组到一个 640g 的显存。那么这是一个什么样的概念？我举一个例子，就是现有的最大的一个异构图的数据集，Open Graph Benchmark 里面最大的一张图是叫 MAG240M，里面是一些论文作者引用网络，大概是 2 亿的点 10 亿的边。
它里面最大的东西是 310G 的特征，就是使用 BERT 抽取的文本特征，每个文本拥有 768 维，这里有大概 300 多 G 特征。但是在这么大的情况下，我们的 A100 80g，如果我们把 8 个卡拼在一起，已经能超过这么一个数据集的容量了。那么我们是不是可以做这样的一个事情？就是 GPU 存储这么大，我们是不是可以把我们整个图算法的流程完全地 GPU 化满足大部分的场景。因此我们也在想怎么对我们图学习的框架进行 GPU 的加速。
02
图学习框架的 GPU 加速
1. 图学习框架的 GPU 加速
说到加速的话，因为大图很难实现。那我们先把概念转到图很小的情况下，我们会怎么做？在图很小的情况下，一个很直接的办法就是我们把图的结构完全的塞到显存里面去。再通过在显存里面实现的一些图采样的算子，那么就可以把我们在图学习算法里面的图存储以及采样的数据流全部的挪到我们的 GPU 上面去，就可以达到很极致的速度。
但是这里最大的问题就是实际上在做的时候，我们的图还是挺大的。如果是大概 200G 的一个图结构，一个单卡的 GPU 已经做不到了。
2. 图学习框架的 GPU 加速：UVA 模式
一个折中的办法就是我们能不能用 GPU 采样的算子访问内存的图数据，做到大图的扩展。这个工作是 2021 年，大概六七月份的时候有个叫 Torch-Quiver 的团队他们做了一个事情，就是把内存当做显存的一块，用一个叫做 UVA 的模式，用 GPU 的采样算子，直接对内存访问去做采样。达到的效果就是会比这纯 CPU 的采样快很多，比 GPU 显存的存储要慢一些，但是图就已经可以变得很大了。
他们具体的做法就是把 GPU 的图结构存到锁页内存里，通过一种叫做统一地址的 Unified Virtual Addressing 技术，把内存当显存，使用 GPU 采样的 Kernel 就可以直接访问统一的一个地址。他们给出了一组数据，就是当用 UVA 的情况下做采样，会相比纯 CPU 的接近快 20 倍的速度。
3. 图学习框架的 GPU 加速：子图特征拉取
除了 UVA 模式下单卡用内存配合，那么如果我们在多卡的情况下，刚刚也说到了我们 8 块 A100 是能达到 640G 的。我们也是在想能不能用到例如 GPU 之间的通讯NVLink 技术，用 GPU 间通讯传输大的带宽，顶替掉内存到显存上面的传输。
这里我们针对图里面最大的一块特征按列切割的形式进行分卡的存储。例如我们左边的这里有一个 300G 的图特征的存储，我们把它切到每一张卡，大概每一张卡是 40G 的显存消耗。那么我们就可以把一整个 MAG240 加上图结构在 8 个 A100 上面存储，组成了一个很大的共享的 Tensor。
在实际的训练里面，我这里用两个 GPU 来举例子，当我们在做数据并行的时候，每个 GPU 只是负责某一部分节点的训练，采样它子图 ID 的时候，我们可以按列用一次的 All-gather，把所有的点对应的一块子图，全部通过 All-gather 获取，然后在局部的按列切分的特征上面把一些 feature 获取。获取之后再通过 All2all 去交换特征，把整个图的结构特征还原。通过这样的 All2all 把整个特征切分，用 GPU 做传输，我们目前可以把 R-GAT 在 MAG240 上面半分钟就可以训完一个 epoch 了，相比之前的相对于 CPU 的或者说 DGL 的那种也是 CPU 模式的速度会更快。
在 2021 年 KDD 比赛里面我们提了一个方法叫 R-UniMP，里面做了很多对于异构图采样的优化，以及对于图里面一些标签应用的优化。当时我们的模型做得特别复杂，当时我们是训练 100 个 epoch 得到 SoTA 成绩是大概要 40 个小时。今年我们通过一系列的优化，在 8 卡的 A100 上 1.1 小时就能出一个去年最好的模型，也给我们今年参加 NEURIPS22 的比赛提供了很大的算法调研的空间。
03
工业级大规模图表示学习 PGLBox
1. 图表示学习在互联网产品的应用
刚刚说到的几个场景都是偏向学术界的一些怎么去训练大图的场景。通常来讲，学术界可能关注的是节点分类，节点分类的数据结构是很整齐的，也不需要大规模的节点特征。
在工业界的场景，特别是互联网的产品应用，我们在做推荐的时候，通常会面临一个事情，就是对我们的用户以及物品去建模。推荐里面有两个常用的算法，一个是 Item-based 协同过滤，就是去衡量两个物品之间的一些相似，以及 User-base 协同过滤，就是去衡量两个用户之间的相似度。
通常图学习在互联网里面的一些应用大部分都是基于这样的情况，就是我们会在海量的日志里面用户以及物品构造一张很大的图，包括有社会的关系、用户的一些行为以及物品的一些关联。我们通过这些图表示学习的一些方法对每个点给一个向量，在这个向量空间里面可以用距离去度量，就是不同的物品以及用户之间的一些相似度。
在传统的一些算法上面，包括像矩阵分解，也是可以做这样的向量表示的，以及像几年前 Deepwalk 和 Meta-Path 这样的算法，就是通过随机游走加 Word2vec 的算法，也可以构造出图的表示。今年来比较火热的就是用图对比学习的方法，通过图的数据增强，通过对比学习对两个子图去抽取用户和物品的行为表示。
2. 作为图学习的框架需要满足业务调研的需求
作为一个图学习框架，最重要的是需要满足我们业务调研的一些需求，我们也整理了目前大部分的图学习的方法，发现基本上是可以抽象为下面四个步骤的。
第一个是大家选用的图可能是不同的，例如 item 的一种 session graph 就是同构图来的，用户物品可能是个二部图。有些像异构图的话，就可能是在同一个图里面有点击关系，喜好关系或者一些购买关系，那它是一种复杂的异构图的形式。
在有了图之后，做一些图上面的自监督度或无监督任务的话，我们通常会用一些随机游走的方法来产生一些训练的样本。
有了样本之后，对于每一个点我们会采用不同的形式做表达，例如我们会做不同邻居的采样，包括像一个比较出名的算法叫 PinSage，就是通过随机游走选邻居，不同的算法会有自己的定义。
在数据构成之后，最后才是我们的图网络的结构，就是我们怎么去选不同的结构去建模邻居和自身的表达。
目前基本上所有的无监督的图学习的方法，都可以抽象成这么四个流程。
3. 用配置完成工业级的图表示学习训练
2021年，在我们 PGL 内部，我们提供了一套叫 Graph4Rec 的一个工具，可以通过五项的配置，就是配置好图数据，配置图的随机游走样本，采样的算法，正负样本采样生成的方法，以及网络选择，就能自动生成一套表示学习的训练，得到每个节点高质量的 embedding。但是之前我们这套方法是基于分布式的 CPU 去训练的。
4. 大规模应用场景下的痛点
在大规模应用场景下，就出现了现在这样的一些痛点，例如我们通常会有一些千亿级的稀疏特征，包括我们用来存一些用户的 ID、图文 ID 以及一些离散化的图文视频的属性。在这种情况下，我们通常需要一套大规模参数服务器，在训练的时候要跟参数服务器进行大量 embedding 的一些查询及交互。
2021年我们也还是用 CPU 参数服务器加 CPU 图引擎去配合一个 CPU 的 MPI集群做训练。它最大的问题在于现阶段不同模态的建模越来越相似，包括用 Transformer 的一些结构的话，是无法去满足一些复杂的模型的。
其次就是当我们的图采样邻居特别多的时候，拉取 embedding 以及图的采样的情况下，它通讯量很大。我们也经常会发现，例如我们搞 20 台机器去训练的时候，只要有一台通讯被打挂了，这个训练就很容易失败。
有人也可能会说就是能不能用 CPU 参数服务器加图引擎，去配合 GPU 去训练，满足刚刚说到的一些复杂模型的需求。但是我们也会发现 IO 是很密集的，就是一些参数服务器的通讯，也是会导致整体的 GPU 利用率很低。
5. PGLBox：全流程 GPU 化的大规模图学习引擎
延续我们最开始说到的，就是既然现在的 GPU 服务器已经例如 8 块 A100 我们有 640G 的一个显存，甚至内存可能都已经有三个 T 的情况下，那么我们是不是就可以做一套全流程完全 GPU 化的大规模图学习引擎？就是把我们的图结构、图的特征表以及我们的参数服务器全部塞到一台机器的 GPU 里面去，即通过 GPU 的分卡，全部塞到我们的 GPU 里面去。那么就把我们刚刚说到的所有的跨机的通讯全部取消掉。在这个情况下，640G 已经能支持很大的一个量级的业务了。所以这里我们也是实现了一个全流程 GPU 优化的一个大规模图学习引擎，也是为了方便我们的一些业务。
因为跟学术界的不一样，就是大家的数据很规整，我们的业务可能每次过来的就是用户给两个 ID，我们就要去构图，或就做查询。所以我们也做了易用性上面的很多个 GPU 的哈希表，去支持哈希的构图以及哈希的图查询，也通过 GPU 加速这种复杂异构图的采样。以及我们在 GPU 上面做了很多的参数服务器的梯度优化策略，来减少显存的使用。
6. PGLBox：通过层次化存储和流水线，支持百亿节点/数百亿边
如果业务再大一点，就是我们的图可能再大到百亿点或者百亿边以上的规模，我们也提供了另外一种层次化的策略，就是把我们的整一个大图存到一个 SSD 以及内存上面去，当我们使用到一个训练的时候，再把对应的东西挪到显存里面去。
这里的图表示学习训练我们做了一系列采样的情况下，通常我们训练的一块它是很集中的，就集中在某一块的子图。所以在这里我们就引入了一个叫 Pass 的一个概念，就是我们一个 epoch 可能要把整个图游走，全部的点给遍历了一遍。但是我们的一个 Pass 就是把整一个 epoch 把它拆成若干个小的 Pass，在 Pass 里面做样本的游走，以及参数服务器的一些特征的准备。准备完了之后，这一块就很适合直接在显存里面做。那么在显存里面我们就一直训了 N 个 Batch，把这整个 Pass 训练完之后，我们再把对应的一些 embedding 以及吸收的一些参数同步到一个 SSD 以及内存的参数服务器上面去。
不同的 Pass 之间，我们可以做一个流水线的并行。就例如我们在做 pull embedding，从内存去拉取对应的一个 embedding 的时候，第二个 Pass 的游走的采样以及样本的生成和分析已经可以做了。所以通过这样的一个 Pass 级的一个训练并行，尽管我们引入了这些层次化的存储，速度也是跟全显存的模式能做到一个持平的效果。
7. PGLBox 加速下的图表示学习
有了这么一个 PGLBox 的大图的表示学习的框架。我们在公司内也做了很多业务的升级。包括像原来有些 Deepwalk 模型，可能是 480 分钟能做完的，现在已经可以一个小时内就解决了。更复杂的模型，像 GraphSAGE 这种的就是会随着我们采样的邻居个数，导致计算量指数上涨的，在子图结构的指数上涨的同时，特征的拉取以及通信量也是在指数上升的。在 GPU 这种单机 8 卡的情况下，就基本上能把这些通讯给优化掉，像这种复杂的 GraphSAGE 一样的图网络，也能去做到 28 倍的提速。
除此之外也是顺应时代潮流，我们通过这种全局 GPU 化能做到一些复杂模型的支持，像我们应用场景里面可能有不同的一些复杂关系，我们的节点本身是多模态的，里面可能有用户、文本和图像，大规模的参数服务器是提供了一些长期行为 ID 的建模，预训练模型就提供了这些跨模态内容的理解。
因为有了这种 GPU 计算的图架构，我们就可以把以往预训练的大模型去跟这些大规模参数服务器放在同一套框架下面训起来，做一个端对端的学习，包括同时去调节预训练大模型里面的一些参数以及大规模的基于 ID 的一些参数服务器。
通过目前这套 PGLBox 的框架，我们实现了单机 8 卡百亿级规模图表示学习训练，相比 CPU 分布式方案提升 28 倍。
8. 飞桨图学习框架 PGL
PGLBox 目前已经在github开源，具体地址如下：
https://github.com/PaddlePaddle/PGL/tree/main/apps/PGLBox
04
问答环节
Q1：GPU 哈希表跟普通的 CPU 哈希表有什么区别？
A1：目前 GPU 哈希表用起来可能没有 CPU 的那么成熟，所以目前那些主流的开源框架，像 DGL 还有 PYG 那些都没有做这个事情。
我们做这个哈希表，它跟 CPU 的也没有太大的区别，主要就是功能上要有这么一个功能，用来简化在 GPU 上做图。
Q2：Cache 方式是提前挪一些到 GPU 上吗？还是全部放到 GPU上？
A2：Cache 方式我们是提前挪一些到 GPU 上。像 MAC240 这种异构图的数据，它全部加起来也就 300 多 G，是因为它每一维可能有 768 维的特征。
那如果没有那么大，例如我们平时做表示学习的 100 维，甚至在做比较早期的推荐系统的一些场景，很多人可能做 8 维，这里可以放很大的一张图了。在这种情况下，图可以做得很大。所以我们基本上就卡到一个程度，一个 Pass 刚好能把显存完全占满，留一部分的空间去做模型的矩阵计算，基本就可以了。
Q3：Pass 之间的流水线并行会导致梯度的过期吗？
A3：目前流水线的话，梯度上面我们会就每一次从 CPU 拉回来的东西是串行的，保证 Pass 去拉 embedding 的时候，下一个 Pass 去拉 embedding 是必须在我这个 Pass 的 push 之后的，所以这里是没有问题的。
具体能并行掉的是游走采样的那个样本，就是样本的构成，以及采样游走样本，这一块也挺占时间的。
Q4：多机的扩展的问题
A4：单机 8 卡的话在卡内就可以用 NVLink 这些技术去做。多机的话还涉及到 embedding，例如跨机里面会包含一些网络带宽的一些问题。
今天的分享就到这里，谢谢大家。
分享嘉宾
INTRODUCTION
黄正杰
百度
资深研发工程师
黄正杰，百度自然语言处理部资深研发工程师，毕业于中山大学长期深耕语义表示计算、图学习等领域的研发工作，目前是飞桨图学习框架 PGL 负责人，并负责相关工作在百度搜索、信息流、商业广告等核心产品的应用落地。在 KDDCUP-2021 图学习比赛获得 2 个赛道冠军，在人工智能顶级会议 KDD、IJCAI 等发表多篇学术论文，拥有多项相关专利。
OpenKG
OpenKG（中文开放知识图谱）旨在推动以中文为核心的知识图谱数据的开放、互联及众包，并促进知识图谱算法、工具及平台的开源开放。
点击阅读原文，进入 OpenKG 网站。