圖神經(jīng)網(wǎng)絡(luò)(Graph Neural Networks, GNNs)在生物醫(yī)藥、知識(shí)圖譜和AI4S等關(guān)鍵領(lǐng)域展現(xiàn)出巨大潛力,隨著真實(shí)應(yīng)用場(chǎng)景中圖規(guī)模迅速擴(kuò)展至億級(jí)邊甚至更高,GNN訓(xùn)練面臨通信開銷與計(jì)算瓶頸等相關(guān)問題。
近日,我中心人工智能部研發(fā)了TAC框架,系統(tǒng)性地融合了數(shù)據(jù)親和緩存填充算法(Affinity-aware Caching Initialization Algorithm,ACI)、稀疏性感知的混合矩陣存儲(chǔ)格式(HybridMatrix)以及多層次細(xì)粒度訓(xùn)練流水線三大關(guān)鍵技術(shù)。TAC 通過提升緩存局部性有效降低全局通信開銷,并借助 Tensor Core 高效加速稀疏矩陣計(jì)算。實(shí)驗(yàn)證明端到端性能相比其他框架顯著提升。
該研究成果已被第31屆ACM SIGPLAN編程語言與并行實(shí)踐原理年會(huì)(PPoPP 2026)錄用。PPoPP是CCF推薦A類會(huì)議,也是高性能計(jì)算與系統(tǒng)領(lǐng)域的頂級(jí)學(xué)術(shù)會(huì)議。本工作得到國(guó)家重點(diǎn)研發(fā)計(jì)劃、中國(guó)科學(xué)院先導(dǎo)專項(xiàng)的支持。論文第一作者為我中心工程師梁智強(qiáng),通訊作者為王玨正高級(jí)工程師和王彥棡研究員,博士研究生高宏宇為共同第一作者。

TAC 框架系統(tǒng)架構(gòu)圖

混合稀疏矩陣格式
相關(guān)成果:
[1] Zhiqiang Liang,Hongyu Gao,Jue Wang,Fang Liu,Xingguo Shi,Junyu Gu,Peng Di,San Li,Lei Tang,Chunbao Zhou,Lian Zhao,Yangang Wang,Xuebin Chi. TAC: Cache-based System for Accelerating Billion-Scale GNN Training on Multi-GPU Platform. ACM SIGPLAN Symposium on Principles & Practice of Parallel Programming(PPoPP’26).
責(zé)任編輯:郎楊琴