PGLBox 超大规模 GPU 端对端图学习训练框架正式发布
百度Geek说约 4 分钟GEEK TALK
01
超高性能 GPU 分布式图学习训练框架
随着图数据规模的不断增大，基于 CPU 分布式的解决方案需要大量的跨机器通讯，导致训练速度慢且稳定性差。为了解决这个问题，PGLBox 将图存储、游走、采样、训练全流程 GPU 化，并实现流水线架构，极致提升异构硬件效率，大幅提升了图学习算法的训练速度。同时，针对 NVLink 拓扑、网卡拓扑非全互联问题，实现智能化中转通信，进一步提升训练能力。相比基于 MPI CPU 分布式的传统方案，训练速度提升 27 倍。PGLBox 实现了多级存储体系，对图、节点属性和图模型进行差异化存储，即图结构全显存、节点属性二级存储和图模型三级存储，将图规模提升了一个数量级。为了平衡磁盘、内存、显存之间的占用，PGLBox 实现了均衡训练，对 Pass 大小平滑处理，削峰填谷，降低峰值显存，使得在单机情况下，可支持的图规模得到大幅提升。
GEEK TALK
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 在百度内大量业务场景实现广泛应用并取得显著业务收益，如广告推荐、信息流推荐、百度搜索、百度网盘、小度平台等。
GEEK TALK
03
在哪里可以找到我们
看到这里相信大家已经迫不及待想要开箱试用了吧！PGLBox 已全面开源，欢迎大家试用或转发推荐，详细代码库链接请戳下方链接！
⭐️ 欢迎STAR收藏 ⭐
https://github.com/PaddlePaddle/PGL/tree/main/apps/PGLBox
更多交流欢迎通过邮件 [email protected] 与我们联系，感谢支持！
END
[1]https://arxiv.org/abs/2112.01035
[2]https://mp.weixin.qq.com/s/aSxFpkyX5MyFYLfZuIagzg
[3]https://ogb.stanford.edu/neurips2022/results/