关于机器学习:阿里云PAIx达摩院GraphScope开源基于PyTorch的GPU加速分布式GNN框架

作者：艾宝乐

导读

近期阿里云机器学习平台 PAI 团队和达摩院 GraphScope 团队联合推出了面向 PyTorch 的 GPU 减速分布式 GNN 框架 GraphLearn-for-PyTorch(GLT)。GLT 利用 GPU 的弱小并行计算性能来减速图采样，并利用 UVA 来缩小顶点和边特色的转换和拷贝。对于大规模图，GLT 应用了生产者 - 消费者的架构，通过异步并发的分布式采样和特色查找以及热点缓存性能反对在多个 GPU 或多个机器上进行高效的分布式训练。接口上，GLT 放弃了 PyTorch 的格调，并且和 PyG 兼容，只需大量代码批改就能够减速 PyG 的单机训练程序，或者将 PyG 单机模型改成分布式训练。此外，GLT 还提供了灵便的分布式训练部署以满足不同的需要。

开源地址：https://github.com/alibaba/graphlearn-for-pytorch

文档地址：https://graphlearn-torch.readthedocs.io/en/latest/index.html

背景介绍

图神经网络作为一种图数据上示意学习的办法曾经被广泛应用在图相干的各个领域，在电商举荐、平安风控、生物分子等畛域获得了理论落地。图神经网络因为其独特的数据处理逻辑和神经网络计算逻辑，须要有专门的学习框架来反对训练。PAI 团队之前开源了大规模工业级分布式图学习框架 GraphLearn（https://github.com/alibaba/graph-learn）。GraphLearn 以 TensorFlow 1.x 系列为主，采纳 ps 架构的异步训练模式，反对十亿节点，百亿边规模的大规模异构图分布式训练，利用于阿里内外部若干业务场景。随着 PyTorch 的风行，其更加灵便的贴近 Python 的接口，简略易调试等个性使得算法开发者更偏向于应用 PyTorch 开发模型。DGL 和 PyG 等基于 PyTorch 的开源 GNN 框架以单机为主，无奈反对大规模图的分布式训练。

此外，因为 GPU 并行计算的劣势，图神经网络应用 GPU 训练比 CPU 训练有数倍的晋升。然而常见的图神经网络框架将图拓扑数据和特色存在内存里，应用 CPU 进行采样和特色查找并将数据拷贝到 GPU 进行神经网络训练，这个过程中图采样和特色查找局部很容易成为整体训练的瓶颈。上面咱们以大规模图上典型的训练流程为例对训练过程的性能瓶颈进行剖析阐明。

一个典型的 GNN 训练流程[1] 包含:

1. 子图拓扑采样，采样多跳街坊并组成子图；
2. 查问子图里节点或者边的特色；
3. 将子图格局转换成神经网络训练须要的格局并且拷贝到 GPU 显存中；
4. 对原始特色进行解决，比方离散特色进行 embedding lookup；
5. 街坊聚合；
6. 节点更新。

其中，3 和 4 为可选步骤。常见的 GNN 模型神经网络参数相对来说比拟小，因而计算量也比拟小，瓶颈通常在 1 - 4 步，具体来说次要是 I / O 操作，包含通信，数据拷贝和格局转换。这导致即便应用 GPU 进行训练，GPU 的资源利用率也很低，整体吞吐以及扩展性很难进步。

综上所述，咱们须要一个高效的基于 PyTorch 的分布式 GNN 训练框架，可能充分利用 GPU 等硬件资源，可能基于图的数据分布性质，联合不同算法模型和并行策略做相应优化，缩小通信和数据转换耗时，晋升整体吞吐。

要害设计

设计初期，咱们和 Quiver[2] 团队单干针对 GPU 采样的可行性进行了初步摸索，发现 GPU 采样相比 CPU 可能带来数量级的晋升。此外，为了缩小特色在 CPU 内存和 GPU 显存之间的拷贝开销，特色能够间接存储在 GPU 显存里。然而对于规模比拟大的数据来说，显存无奈存储所有特色数据，因而咱们用 UVA 技术来进行特色的存储，相比间接内存存储也有数量级的性能晋升。工业界的图规模很容易冲破单机的极限，因而咱们进一步设计了分布式训练框架。具体来说，咱们应用生产者 - 消费者范式来解决分布式图采样和特色查找以及模型训练的关系，将图采样、特色查找与训练进行解耦，应用多过程并行和协程异步并发来减速采样和特色查找，并应用预取和热点缓存的形式进一步缩小训练端的期待，晋升端到端吞吐。思考到用户迁徙老本和易用性，在接口上咱们放弃了和 PyG 的兼容，只需大量改变 PyG 代码就能够减速 PyG 的训练，或者将其迁徙到 GLT 的分布式训练上。以下为咱们具体论述几个要害的设计点。

GPU 采样

GLT 将图拓扑应用 CSR 格局存储在 GPU 显存或者 pin memory 里，实现了 CUDA 采样算子来进行 GPU 并行采样。应用 CSR 存储能够很容易失去每个节点的街坊，并且独立对每个节点进行采样，因而能够不便地利用 GPU 多线程进行并行采样。咱们应用了蓄水池算法来进行无放回随机采样。在 batch size 大的状况下，GPU 比 CPU 采样能有数量级的晋升。

UnifiedTensor

为了打消 CPU 内存到 GPU 显存的拷贝开销，一个比拟直观的办法是将特色寄存在 GPU 显存里，然而因为单卡的显存无限，在特色数据比拟大的状况下也很难齐全把特色存到显存里。因而，在 GLT 中咱们利用图本身的个性如 power law 散布和采样拜访个性如有些度比拟高的节点被拜访的概率高，将局部热点特色寄存在 GPU 显存里，其余特色存在内存，同时须要利用 UVA 让 GPU 拜访内存里的特色。GLT 设计了 UnifiedTensor 将 CUDA Tensor 和 CPU Tensor 对立治理起来，以提供简洁高效的数据拜访。进一步，如果 GPU 之间能够间接进行 peer2peer 拜访（具备 NVLink），这些 GPU 的显存也能够被对立治理，从而扩充特色在显存的存储。GLT 应用 UnifiedTensor 将这些不同硬件设施上的存储对立治理起来，提供间接拜访 CUDA Tensor，通过 NVLink 拜访其余 GPU 上的 CUDA Tensor，并通过 UVA 进行 ZERO-COPY 拜访 CPU Tensor 的能力，下层查找元素接口就像一般 Tensor 一样，底层会主动去对应的设施上进行访存操作。

Feature

Feature 由 UnifiedTensor 形成，具备硬件拓扑感知性能。具体来说，首先，依照用户指定 CPU/GPU 内存大小，对特色进行划分，分为 GPU(hot) 局部和 CPU(cold) 局部。其次，对 GPU 局部，依据用户指定的 replica 策略，进行 replica，包含每个卡 replica 和每个 NVLink 连贯的 GPU group 之间的 replica。GPU group replica 的形式，相比卡间 replica, 能够有更多的 hot data 存在 GPU 上，因为 GPU group 里 GPU 之间都是能够 p2p 拜访的。实现上 GLT 形象出 DeviceGroup 来对立示意卡间 replica 和组间 replica。一个 DeviceGroup 示意一组 NVLink 连贯的 GPUs。假如 8 卡没有 NVLink，那么会对应 8 个 DeviceGroup，如果 GPU 0-3 两两 NVLink 连贯，GPU 4-7 两两 NVLink 连贯，那么 GPU 0-3 为一组 DeviceGroup, GPU 4-7 为一组 DeviceGroup。理论测试中，应用 UnifiedTensor 的 Feature 性能比 CPU Tensor 的查找(包含拷贝到 GPU) 快 1 个数量级，而且能够通过管制 GPU 存储局部的比例来灵便达到速度和显存占用的均衡。

分布式设计

GLT 分布式 GNN 训练次要分成：分布式采样，特色查找，模型训练 3 局部。一次采样的后果个别比拟小(最大为十几 MB)，特色查找的后果比拟大(百 MB)，训练时应用特色查找的后果进行神经网络计算。对特色查找来说须要思考缩小和训练任务之间的数据转换和拷贝。采样和训练之间是典型的生产者和消费者关系，因而能够分成不同工作，通过缓冲区连贯，均衡生产者和消费者的解决能力，起到一个数据缓存的作用，同时也达到了一个解耦的作用。基于生产者 - 消费者形式，GLT 的分布式训练有两种根本类型的过程：采样过程和训练过程。

采样过程：负责分布式街坊采样和特色收集。采样后果将被发送到采样音讯通道，该通道将进一步用于训练任务。

训练过程：对应于 PyTorch DDP 的分布式训练过程，通常，每个训练过程将占用一个 GPU 进行训练。

这些过程能够灵便地散布在不同的机器上，为了更好地治理分布式过程部署，GLT 的分布式训练提供了两种参考部署模式：Worker 模式和 Server-Client 模式。

Worker 模式 里，数据切分后，每个机器持有一个分片，采样过程和训练过程一起部署在这些机器上。每个训练过程能够 spawn 出多个采样子过程，采样过程通过一个共享内存的音讯通道将采样后果传递给训练过程。对于采样过程来说，能够应用多过程进行采样，并且每个分布式采样算子都应用 Python 协程来并发执行，将后果放到音讯通道里。为了缩小音讯通道到训练过程 GPU 的拷贝耗时，音讯通道也能够放到 pin memory 上。

Server-Client 模式 下，集群中存在两种类型的机器节点，即 Server 节点和 Client 节点。采样过程部署在 Server 节点，训练过程散布在所有 Client 节点上。采样过程生成的样本后果将通过一个 RPC 实现的近程音讯通道发送到以后训练过程进行训练。Server-Client 模式能够将采样和训练不同 workload 的工作放到不同机器，进行资源上的解耦。

总体架构

GLT 的次要指标是充分利用 GPU/NVLink/RDMA 等硬件资源和 GNN 模型的个性，减速单机和分布式环境下的端到端 GNN 训练。

存储：在 GPU 训练场景，图采样和 CPU-GPU 数据传输很容易成为次要性能瓶颈。为了减速图采样和特色查找，GLT 实现了 UnifiedTensor 对立 CPU 和 GPU 的内存治理。为了缩小特色收集引起的 CPU-GPU 数据传输开销，GLT 反对将热点特色缓存在 GPU 显存中，并通过 UVA 拜访其余特色数据。咱们还利用高速 NVLink 在 GPU 之间扩大 GPU 缓存的容量。

图操作算子：存储之上，GLT 实现了包含街坊采样、负采样、特色查找、子图采样等同时反对 CPU 和 GPU 图操作算子。

分布式算子：对于分布式训练，为避免近程数据拜访阻塞模型训练过程，GLT 在 PyTorch RPC 之上封装了一个高效的 RPC 框架，并采纳多过程并行和异步并发的图采样和特色查找操作来暗藏网络提早并进步端到端训练吞吐量。

接口 ：为了升高 PyG 用户的学习老本，GLT 的下层 API，如Sampler, Dataset 和Loader，接口上都与 PyG 兼容。因而，PyG 用户只需批改很少的代码即可充分利用 GLT 的减速能力。

模型：因为 GLT 与 PyG 兼容，你能够应用简直任何 PyG 的模型作为根底模型，此外咱们也提供了丰盛的分布式训练示例。

零碎性能

咱们在一台装备 A100 GPU 的机器进行单机扩展性测试，测试环境为 CUDA 11.4、PyTorch 1.12 和 GLT 0.2.0rc2，下图展现了街坊采样和特色查找的总吞吐量。能够看出 GLT 有线性的扩展性(因为有 NVLink，多卡的缓存容量更多，因而会存在超线性减速)。

此外，咱们还测试了多机的分布式采样和特色查找的扩展性。下图展现了每个机器装备 2 个 A100 GPU 的环境下，2 个机器和 4 个机器相比单个机器的吞吐量减速比。测试应用 CUDA11.4、PyTorch 1.12 和 GLT 0.2.0rc2 进行。能够看出，2 机到 4 机也有近线性的扩展性。

最初，咱们测试了分布式 e2e 的性能。咱们在 2 机每机 2 卡 A100 的设置下和 DGL 做了初步比照(DGL 版本 0.9.1，GLT 版本 0.2.0).

结语

本文介绍了基于 PyTorch 的 GPU 减速分布式 GNN 框架 GraphLearn-for-PyTorch(GLT)，GLT 提供了分布式 GPU 训练的优化减速能力，可能充分利用 GPU 等硬件资源进行图采样和特色查找等操作，具备线性扩展性。下层接口上和 PyG 兼容，能够很容易地减速 PyG 已有模型或者将已有模型改成分布式版本。GLT 曾经开源并且在 PyG, GraphScope 中都有示例，前面咱们会继续开发优化，欢送应用和奉献！

[1] P3: Distributed Deep Graph Learning at Scale

[2] Quiver: Supporting GPUs for Low-Latency, High-Throughput GNN Serving with Workload Awareness

欢送各位感兴趣的开发者们退出咱们，与咱们一起共建~

GLT 开源我的项目钉群

收费支付 交互式建模 PAI-DSW、模型训练 PAI-DLC 5000CU* H 计算资源包，以及价值 500 元模型在线服务 PAI-EAS 抵扣包。