近年來，圖神經(jīng)網(wǎng)絡(luò)（GNN）在許多場景中得到廣泛應(yīng)用，如社交網(wǎng)絡(luò)分析、推薦系統(tǒng)、自動駕駛等。然而，真實世界中龐大的數(shù)據(jù)量及圖結(jié)構(gòu)顯著的不規(guī)則稀疏性，為圖神經(jīng)網(wǎng)絡(luò)的低時延、高能效應(yīng)用帶來了巨大的挑戰(zhàn)。隨著社會圖數(shù)據(jù)規(guī)模的爆炸式增長，解決圖神經(jīng)網(wǎng)絡(luò)中超大規(guī)模數(shù)據(jù)量帶來的訪存及計算問題變得越發(fā)迫切。

中國科學(xué)院自動化研究所程健課題組聯(lián)合上海交通大學(xué)先進體系結(jié)構(gòu)實驗室的梁曉峣教授和李鋼博士從圖神經(jīng)網(wǎng)絡(luò)壓縮算法設(shè)計、專用推理架構(gòu)設(shè)計及高效訓(xùn)練框架設(shè)計三方面對圖神經(jīng)網(wǎng)絡(luò)的高效計算進行了系統(tǒng)和深入地研究，取得系列進展。相關(guān)成果相繼發(fā)表于第11屆國際表征學(xué)習(xí)大會（ICLR 2023）、第30屆IEEE國際高性能計算機體系結(jié)構(gòu)大會（HPCA 2024）和第29屆ACM架構(gòu)支持編程語言和操作系統(tǒng)國際會議（ASPLOS 2024），第一作者均為自動化所博士生朱澤雨。

圖1. 圖數(shù)據(jù)具有復(fù)雜的拓撲結(jié)構(gòu)

在算法設(shè)計方面，研發(fā)團隊提出了端到端的圖拓撲自適應(yīng)的混合量化算法A2Q，在保證模型精度的前提下極大地壓縮了圖數(shù)據(jù)。實驗結(jié)果表明，相比于使用FP32格式存儲的圖神經(jīng)網(wǎng)絡(luò)模型，該壓縮方法能夠在精度幾乎無損的情況下實現(xiàn)高達18.6倍的壓縮比。

圖2. 研究發(fā)現(xiàn)聚合后節(jié)點特征的數(shù)值范圍與其拓撲特性強相關(guān)，基于此特性提出了端到端的圖拓撲自適應(yīng)的混合量化算法A2Q

A2Q較高的模型壓縮率在降低系統(tǒng)計算延遲和能耗方面展示出了巨大的潛力，然而細粒度的混合精度數(shù)據(jù)表示和無規(guī)則稀疏的節(jié)點特征使得該算法在通用處理器上難以充分發(fā)揮優(yōu)勢。

為此，在專用推理架構(gòu)方面，團隊針對壓縮算法特性和圖數(shù)據(jù)特性設(shè)計了一款專用加速器MEGA。通過比特位自適應(yīng)的存儲方式和高效的在線編解碼壓縮方法，最大限度減少了存儲碎片化，降低了數(shù)據(jù)傳輸開銷。此外，MEGA采用的Condense-Edge調(diào)度執(zhí)行方式，能夠?qū)⑾∈璧倪呥B接聚合在一起以增強訪問節(jié)點特征時的空間局部性，提升數(shù)據(jù)復(fù)用度。在與英偉達 RTX3090 GPU相當?shù)姆逯邓懔υO(shè)定下，MEGA實現(xiàn)了121倍的推理加速及476倍的能效提升。

圖3. MEGA加速器架構(gòu)圖

更進一步，研發(fā)團隊提出了全流程圖神經(jīng)網(wǎng)絡(luò)訓(xùn)練加速系統(tǒng)FastGL，極大地提升了圖神經(jīng)網(wǎng)絡(luò)在超大規(guī)模圖數(shù)據(jù)上的訓(xùn)練速度。與最先進的圖神經(jīng)網(wǎng)絡(luò)訓(xùn)練加速系統(tǒng)DGL對比，F(xiàn)astGL在多種超大規(guī)模圖數(shù)據(jù)和圖神經(jīng)網(wǎng)絡(luò)模型上實現(xiàn)了平均2.2倍的加速效果，并且具有更好的擴展性，有利于應(yīng)對未來圖數(shù)據(jù)規(guī)模持續(xù)增加的趨勢。

圖4. 全流程圖神經(jīng)網(wǎng)絡(luò)訓(xùn)練加速系統(tǒng)FastGL訓(xùn)練框架示意圖

這一系列圖神經(jīng)網(wǎng)絡(luò)大規(guī)模圖數(shù)據(jù)推理及訓(xùn)練加速方面的研究成果貫穿算法、架構(gòu)和GPU高效訓(xùn)練框架，形成了較為系統(tǒng)的圖神經(jīng)網(wǎng)絡(luò)高效計算方案，也為解決通用大模型推理中由超大規(guī)模參數(shù)量和超長KV cache引起的計算、訪存雙密集加速等難題提供了啟發(fā)。

發(fā)表該系列研究的國際會議HPCA、ASPLOS與ISCA、MICRO并稱為計算機體系結(jié)構(gòu)領(lǐng)域的四大頂會，是計算機體系結(jié)構(gòu)研究的風(fēng)向標，ICLR則為深度學(xué)習(xí)領(lǐng)域的頂級學(xué)術(shù)會議。

論文信息：

[1] Zeyu Zhu, Fanrong Li, Zitao Mo, Qinghao Hu, Gang Li, Zejian Liu, Xiaoyao Liang, Jian Cheng. A2Q: Aggregation-Aware Quantization for Graph Neural Networks. ICLR, 2023

[2] Zeyu Zhu*, Fanrong Li*, Gang Li, Zejian Liu, Zitao Mo, Qinghao Hu, Xiaoyao Liang, Jian Cheng. MEGA: A Memory-Efficient GNN Accelerator Exploiting Degree-Aware Mixed-Precision Quantization. HPCA, 2024

[3] Zeyu Zhu, Peisong Wang, Qinghao Hu, Gang Li, Xiaoyao Liang, Jian Cheng. FastGL: A GPU-Efficient Framework for Accelerating Sampling-Based GNN Training at Large Scale. ASPLOS, 2024

相關(guān)開源代碼：

A2Q: https://github.com/weihai-98/A-2Q

FastGL: https://github.com/a1bc2def6g/fastgl-ae

來源：中國科學(xué)院自動化研究所