圖神經網絡能夠處理復雜的圖結構數(shù)據(jù),在芯片設計、能源規(guī)劃、推薦系統(tǒng)等多個領域具有廣泛應用前景。隨著數(shù)據(jù)規(guī)模和復雜性的增長,單個計算單元已無法滿足大規(guī)模圖數(shù)據(jù)的處理需求。盡管使用多GPU加速GNN訓練逐漸成為主流,全批次圖神經網絡的分布式訓練仍面臨嚴重負載不平衡和高通信開銷等挑戰(zhàn)。
我中心人工智能技術與應用發(fā)展部門設計了高效全批次訓練系統(tǒng)ParGNN。該系統(tǒng)采用了一種由負載均衡器指導的自適應負載均衡方法,并結合圖過劃分策略有效緩解了計算負載不平衡的問題。ParGNN實現(xiàn)了新型子圖流水線算法,使得計算和通信過程得以重疊,從而在不影響GNN訓練精度的前提下大幅提高了效率。實驗表明,相較于當前最先進的解決方案DGL和Pipe GCN,ParGNN不僅能實現(xiàn)最高的訓練精度,還能在最短時間內達到預設目標精度。此項研究成果已被第62屆國際設計自動化會議Design Automation Conference(DAC 2025,CCF推薦A類會議)錄用。該成果得到國家重點研發(fā)計劃(2023YFB4502303)和中國科學院先導專項(XDB0500103)的支持。
ParGNN總體設計
論文第一作者為我中心博士研究生顧峻瑜,我中心副研究員曹榮強和博士研究生李順德為共同一作,通信作者為我中心王玨正高級工程師。
相關成果
Junyu Gu,Shunde Li,Rongqiang Cao,Jue Wang,Zijian Wang,Zhiqiang Liang,Fang Liu,Shigang Li,Chunbao Zhou,Yangang Wang,Xuebin Chi. ParGNN: A Scalable Graph Neural Network Training Framework on multi-GPUs. Proceedings of the 62st ACM/IEEE Design Automation Conference. 2025.
責任編輯:郎楊琴