PGLBox 超大規模 GPU 端對端圖學習訓練框架正式發佈

語言: CN / TW / HK

作者 | 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 在百度內大量業務場景實現廣泛應用並取得顯著業務收益,如廣告推薦、信息流推薦、百度搜索、百度網盤、小度平台等。

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/

推薦閲讀:

百度工程師帶你瞭解Module Federation

巧用Golang泛型,簡化代碼編寫

Go語言DDD實戰初級篇

Diffie-Hellman密鑰協商算法探究

貼吧低代碼高性能規則引擎設計

淺談權限系統在多利熊業務應用