阿里云PAIx达摩院GraphScope开源基于PyTorch的GPU加速分布式GNN框架

本文涉及的产品
交互式建模 PAI-DSW,每月250计算时 3个月
模型在线服务 PAI-EAS,A10/V100等 500元 1个月
模型训练 PAI-DLC,100CU*H 3个月
简介: 阿里云机器学习平台 PAI 团队和达摩院 GraphScope 团队联合推出了面向 PyTorch 的 GPU 加速分布式 GNN 框架 GraphLearn-for-PyTorch(GLT) 。

作者:艾宝乐

导读

近期阿里云机器学习平台 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 一样,底层会自动去对应的设备上进行访存操作。

1.jpg

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 存储部分的比例来灵活达到速度和显存占用的平衡。

2.jpg

3.jpg

分布式设计

GLT分布式GNN训练主要分成:分布式采样,特征查找, 模型训练3部分。一次采样的结果一般比较小(最大为十几MB),特征查找的结果比较大(百MB),训练时使用特征查找的结果进行神经网络计算。对特征查找来说需要考虑减少和训练任务之间的数据转换和拷贝。采样和训练之间是典型的生产者和消费者关系,因此可以分成不同任务,通过缓冲区连接,平衡生产者和消费者的处理能力,起到一个数据缓存的作用,同时也达到了一个解耦的作用。基于生产者-消费者方式,GLT的分布式训练有两种基本类型的进程:采样进程和训练进程。

采样进程:负责分布式邻居采样和特征收集。采样结果将被发送到采样消息通道,该通道将进一步用于训练任务。

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

这些进程可以灵活地分布在不同的机器上,为了更好地管理分布式进程部署,GLT的分布式训练提供了两种参考部署模式:Worker 模式和 Server-Client 模式。

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

4.jpg

Server-Client模式下,集群中存在两种类型的机器节点,即 Server 节点和 Client 节点。采样进程部署在 Server 节点,训练进程分布在所有 Client 节点上。采样进程生成的样本结果将通过一个RPC实现的远程消息通道发送到当前训练进程进行训练。Server-Client 模式可以将采样和训练不同 workload 的任务放到不同机器,进行资源上的解耦。

5.jpg

总体架构

6.jpg

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, DatasetLoader,接口上都与PyG兼容。因此,PyG 用户只需修改很少的代码即可充分利用 GLT 的加速能力。

模型:由于 GLT 与 PyG 兼容,你可以使用几乎任何 PyG 的模型作为基础模型,此外我们也提供了丰富的分布式训练示例。

系统性能

我们在一台配备A100 GPU的机器进行单机扩展性测试,测试环境为 CUDA 11.4、PyTorch 1.12 和 GLT 0.2.0rc2,下图展示了邻居采样和特征查找的总吞吐量。可以看出 GLT 有线性的扩展性(由于有NVLink,多卡的缓存容量更多,因此会存在超线性加速)。

7.jpg

此外,我们还测试了多机的分布式采样和特征查找的扩展性。下图展示了每个机器配备2个A100 GPU的环境下,2个机器和4个机器相比单个机器的吞吐量加速比。测试使用 CUDA11.4、PyTorch 1.12和 GLT 0.2.0rc2 进行。可以看出,2机到4机也有近线性的扩展性。

8.jpg

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

9.jpg

结语

本文介绍了基于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

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

相关实践学习
部署Stable Diffusion玩转AI绘画(GPU云服务器)
本实验通过在ECS上从零开始部署Stable Diffusion来进行AI绘画创作,开启AIGC盲盒。
相关文章
|
5天前
|
存储 机器学习/深度学习 人工智能
2025年阿里云GPU服务器租用价格、选型策略与应用场景详解
随着AI与高性能计算需求的增长,阿里云提供了多种GPU实例,如NVIDIA V100、A10、T4等,适配不同场景。2025年重点实例中,V100实例GN6v单月3830元起,适合大规模训练;A10实例GN7i单月3213.99元起,适用于混合负载。计费模式有按量付费和包年包月,后者成本更低。针对AI训练、图形渲染及轻量级推理等场景,推荐不同配置以优化成本和性能。阿里云还提供抢占式实例、ESSD云盘等资源优化策略,支持eRDMA网络加速和倚天ARM架构,助力企业在2025年实现智能计算的效率与成本最优平衡。 (该简介为原文内容的高度概括,符合要求的字符限制。)
|
2天前
|
机器学习/深度学习 人工智能 并行计算
Unsloth:学生党福音!开源神器让大模型训练提速10倍:单GPU跑Llama3,5小时变30分钟
Unsloth 是一款开源的大语言模型微调工具,支持 Llama-3、Mistral、Phi-4 等主流 LLM,通过优化计算步骤和手写 GPU 内核,显著提升训练速度并减少内存使用。
115 3
Unsloth:学生党福音!开源神器让大模型训练提速10倍:单GPU跑Llama3,5小时变30分钟
|
10天前
|
存储 机器学习/深度学习 PyTorch
PyTorch Profiler 性能优化示例:定位 TorchMetrics 收集瓶颈,提高 GPU 利用率
本文探讨了机器学习项目中指标收集对训练性能的影响,特别是如何通过简单实现引入不必要的CPU-GPU同步事件,导致训练时间增加约10%。使用TorchMetrics库和PyTorch Profiler工具,文章详细分析了性能瓶颈的根源,并提出了多项优化措施
31 1
PyTorch Profiler 性能优化示例:定位 TorchMetrics 收集瓶颈,提高 GPU 利用率
|
5天前
|
机器学习/深度学习 存储 人工智能
2025年阿里云GPU服务器的租赁价格与选型指南
随着AI、深度学习等领域的发展,GPU服务器成为企业及科研机构的核心算力选择。阿里云提供多种GPU实例类型(如NVIDIA V100、A100等),涵盖计算型、共享型和弹性裸金属等,满足不同场景需求。本文详解2025年阿里云GPU服务器的核心配置、价格策略及适用场景,帮助用户优化选型与成本控制,实现高效智能计算。
|
5天前
|
机器学习/深度学习 存储 弹性计算
阿里云gpu云服务器租用价格:最新收费标准及活动价格参考
阿里云gpu云服务器多少钱?A10卡GN7i GPU云服务器32核188G3213.99/1个月起,V100卡GN6v GPU云服务器8核32G3830.00/1个月起,阿里云GPU云服务器是基于GPU应用的计算服务,多适用于视频解码,图形渲染,深度学习,科学计算等应用场景,该产品具有超强计算能力、网络性能出色、购买方式灵活、高性能实例存储( GA1和GN5特有)等特点。下面小编来介绍下阿里云gpu云服务器最新的收费标准及活动价格。
|
1月前
|
人工智能 文字识别 异构计算
NVIDIA-Ingest:英伟达开源智能文档提取及结构化工具,支持 GPU 加速和并行处理
NVIDIA-Ingest 是英伟达开源的智能文档提取工具,支持 PDF、Word、PPT 等多种格式,提供并行处理和 GPU 加速,适用于企业内容管理和生成式应用。
114 18
NVIDIA-Ingest:英伟达开源智能文档提取及结构化工具,支持 GPU 加速和并行处理
|
24天前
|
机器学习/深度学习 人工智能 弹性计算
阿里云AI服务器价格表_GPU服务器租赁费用_AI人工智能高性能计算推理
阿里云AI服务器提供多种配置,包括CPU+GPU、FPGA等,适用于人工智能、机器学习和深度学习等计算密集型任务。本文整理了阿里云GPU服务器的优惠价格,涵盖NVIDIA A10、V100、T4等型号,提供1个月、1年和1小时的收费明细。具体规格如A10卡GN7i、V100-16G卡GN6v等,适用于不同业务场景,详情见官方页面。
124 11
|
2月前
|
机器学习/深度学习 人工智能 PyTorch
阿里云GPU云服务器怎么样?产品优势、应用场景介绍与最新活动价格参考
阿里云GPU云服务器怎么样?阿里云GPU结合了GPU计算力与CPU计算力,主要应用于于深度学习、科学计算、图形可视化、视频处理多种应用场景,本文为您详细介绍阿里云GPU云服务器产品优势、应用场景以及最新活动价格。
阿里云GPU云服务器怎么样?产品优势、应用场景介绍与最新活动价格参考
|
2天前
|
机器学习/深度学习 PyTorch TensorFlow
深度学习工具和框架详细指南:PyTorch、TensorFlow、Keras
在深度学习的世界中,PyTorch、TensorFlow和Keras是最受欢迎的工具和框架,它们为研究者和开发者提供了强大且易于使用的接口。在本文中,我们将深入探索这三个框架,涵盖如何用它们实现经典深度学习模型,并通过代码实例详细讲解这些工具的使用方法。
|
1月前
|
机器学习/深度学习 并行计算 PyTorch
TorchOptimizer:基于贝叶斯优化的PyTorch Lightning超参数调优框架
TorchOptimizer 是一个基于贝叶斯优化方法的超参数优化框架,专为 PyTorch Lightning 模型设计。它通过高斯过程建模目标函数,实现智能化的超参数组合选择,并利用并行计算加速优化过程。该框架支持自定义约束条件、日志记录和检查点机制,显著提升模型性能,适用于各种规模的深度学习项目。相比传统方法,TorchOptimizer 能更高效地确定最优超参数配置。
131 7