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