PGLBox 超大规模 GPU 端对端图学习训练框架正式发布

news2024/9/22 12:20:56

在这里插入图片描述

作者 | PGLBox项目组

导读

PGLBox是百度研发的基于GPU的大规模图模型训练框架，支持数百亿节点和边的图模型全GPU训练，已在百度广泛部署。相比业界主流的分布式 CPU 解决方案，PGLBox 具有超高性能、超大规模、算法丰富、灵活易用、落地广泛等优势。与传统的基于GPU的分布式解决方案相比，PGLBox可以在相同的成本下提升27倍的训练速度。

全文2232字，预计阅读时间6分钟。

图神经网络（Graph Neural Network，GNN）是近年来出现的一种利用深度学习直接对图结构数据进行学习的方法。通过在图的节点和边上制定聚合的策略，GNN 能够学习到图结构数据中节点以及边内在规律和更加深层次的语义特征。图神经网络不仅成为学术界研究热点，而且已经在工业界广泛应用落地。特别在搜索、推荐、地图等领域，采用大规模分布式图引擎对异构图结构进行建模，这已经成为技术发展的新趋势。目前，分布式图学习框架通常在 CPU 集群上部署分布式图服务以及参数服务器，来支持大规模图结构的存储以及特征的更新。然而，基于 CPU 算力的图学习框架在建设成本、训练速度、稳定性以及复杂算法支持等方面都存在不足。

因此，百度飞桨推出了**能够同时支持复杂图学习算法+超大图+超大离散模型的 GPU 大规模图学习训练框架 PGLBox。**该框架结合了百度移动生态模型团队在大规模业务技术的深耕，凝聚飞桨图学习 PGL 丰富的算法能力与应用经验，并依托飞桨深度学习平台通用的训练框架能力与灵活组网能力。不仅继承了飞桨前期开源的 Graph4Rec[1]的超大规模、灵活易用和适用性广的优点[2]，而且训练性能获得了显著提升，图算法能力支持更广泛。

01 超高性能 GPU 分布式图学习训练框架

随着图数据规模的不断增大，基于 CPU 分布式的解决方案需要大量的跨机器通讯，导致训练速度慢且稳定性差。为了解决这个问题，PGLBox 将图存储、游走、采样、训练全流程 GPU 化，并实现流水线架构，极致提升异构硬件效率，大幅提升了图学习算法的训练速度。同时，针对 NVLink 拓扑、网卡拓扑非全互联问题，实现智能化中转通信，进一步提升训练能力。相比基于 MPI CPU 分布式的传统方案，训练速度提升 27 倍。PGLBox 实现了多级存储体系，对图、节点属性和图模型进行差异化存储，即图结构全显存、节点属性二级存储和图模型三级存储，将图规模提升了一个数量级。为了平衡磁盘、内存、显存之间的占用，PGLBox 实现了均衡训练，对 Pass 大小平滑处理，削峰填谷，降低峰值显存，使得在单机情况下，可支持的图规模得到大幅提升。

02 全面升级预置的图表示学习算法

图节点的属性是多种多样的，可以是文本、图像，也可以是用户画像、地理位置等，如何更好地建模节点特征是图表示学习的一个重要挑战。随着预训练模型席卷 NLP、CV等领域，预训练 Transformer 是节点属性建模不可或缺的一部分。而复杂结构的 Transformer 等预训练模型的引入所增加的大量计算量，是以往 CPU 分布式图表示学习框架不可接受的。**得益于 PGLBox 同时兼备 GPU 的运算能力和大规模图的支持，让我们能够同时实现大规模预训练模型 + 大规模图结构信息 + 大规模离散特征的端对端统一建模。**在大规模图数据，通过三级存储加载之后，我们可以通过加载不同的大规模预训练模型（例如 ERNIE 语言大模型、ERNIE-ViL 跨模态大模型等）来建模更丰富的节点信息。对于大规模离散特征如用户ID、商品ID等，我们可以同时利用到 PGLBox 提供的 GPU 参数服务器能力来建模。最后通过图信息汇聚的 Graph Transformer 图神经网络模块完成信息聚合，得到图的最终表示，并配合下游任务实现跨模态异构图端对端优化。

基于 PGLBox 的 GNN 技术获得了 NeurIPS 2022 大规模 GNN 挑战赛冠军[3]，同时入选了百度 Create2022 十大黑科技，并在 WAVE SUMMIT 2022 上作为飞桨 2.4 版本最重要的框架新特性之一发布。凭借其超高性能、超大规模、超强图学习算法、灵活易用等特性，PGLBox 在百度内大量业务场景实现广泛应用并取得显著业务收益，如广告推荐、信息流推荐、百度搜索、百度网盘、小度平台等。