近年來(lái)，圖神經(jīng)網(wǎng)絡(luò)（GNN）在許多場(chǎng)景中得到廣泛應(yīng)用，如社交網(wǎng)絡(luò)分析、推薦系統(tǒng)、自動(dòng)駕駛等。然而，真實(shí)世界中龐大的數(shù)據(jù)量及圖結(jié)構(gòu)顯著的不規(guī)則稀疏性，為圖神經(jīng)網(wǎng)絡(luò)的低時(shí)延、高能效應(yīng)用帶來(lái)了巨大的挑戰(zhàn)。隨著社會(huì)圖數(shù)據(jù)規(guī)模的爆炸式增長(zhǎng)，解決圖神經(jīng)網(wǎng)絡(luò)中超大規(guī)模數(shù)據(jù)量帶來(lái)的訪存及計(jì)算問(wèn)題變得越發(fā)迫切。
中國(guó)科學(xué)院自動(dòng)化研究所程健課題組聯(lián)合上海交通大學(xué)先進(jìn)體系結(jié)構(gòu)實(shí)驗(yàn)室的梁曉峣教授和李鋼博士從圖神經(jīng)網(wǎng)絡(luò)壓縮算法設(shè)計(jì)、專(zhuān)用推理架構(gòu)設(shè)計(jì)及高效訓(xùn)練框架設(shè)計(jì)三方面對(duì)圖神經(jīng)網(wǎng)絡(luò)的高效計(jì)算進(jìn)行了系統(tǒng)和深入地研究，取得系列進(jìn)展。相關(guān)成果相繼發(fā)表于第11屆國(guó)際表征學(xué)習(xí)大會(huì)（ICLR 2023）、第30屆IEEE國(guó)際高性能計(jì)算機(jī)體系結(jié)構(gòu)大會(huì)（HPCA 2024）和第29屆ACM架構(gòu)支持編程語(yǔ)言和操作系統(tǒng)國(guó)際會(huì)議（ASPLOS 2024），第一作者均為自動(dòng)化所博士生朱澤雨。
圖1. 圖數(shù)據(jù)具有復(fù)雜的拓?fù)浣Y(jié)構(gòu)
在算法設(shè)計(jì)方面，研發(fā)團(tuán)隊(duì)提出了端到端的圖拓?fù)渥赃m應(yīng)的混合量化算法A2Q，在保證模型精度的前提下極大地壓縮了圖數(shù)據(jù)。實(shí)驗(yàn)結(jié)果表明，相比于使用FP32格式存儲(chǔ)的圖神經(jīng)網(wǎng)絡(luò)模型，該壓縮方法能夠在精度幾乎無(wú)損的情況下實(shí)現(xiàn)高達(dá)18.6倍的壓縮比。
圖2. 研究發(fā)現(xiàn)聚合后節(jié)點(diǎn)特征的數(shù)值范圍與其拓?fù)涮匦詮?qiáng)相關(guān)，基于此特性提出了端到端的圖拓?fù)渥赃m應(yīng)的混合量化算法A2Q
A2Q較高的模型壓縮率在降低系統(tǒng)計(jì)算延遲和能耗方面展示出了巨大的潛力，然而細(xì)粒度的混合精度數(shù)據(jù)表示和無(wú)規(guī)則稀疏的節(jié)點(diǎn)特征使得該算法在通用處理器上難以充分發(fā)揮優(yōu)勢(shì)。
為此，在專(zhuān)用推理架構(gòu)方面，團(tuán)隊(duì)針對(duì)壓縮算法特性和圖數(shù)據(jù)特性設(shè)計(jì)了一款專(zhuān)用加速器MEGA。通過(guò)比特位自適應(yīng)的存儲(chǔ)方式和高效的在線(xiàn)編解碼壓縮方法，最大限度減少了存儲(chǔ)碎片化，降低了數(shù)據(jù)傳輸開(kāi)銷(xiāo)。此外，MEGA采用的Condense-Edge調(diào)度執(zhí)行方式，能夠?qū)⑾∈璧倪呥B接聚合在一起以增強(qiáng)訪問(wèn)節(jié)點(diǎn)特征時(shí)的空間局部性，提升數(shù)據(jù)復(fù)用度。在與英偉達(dá) RTX3090 GPU相當(dāng)?shù)姆逯邓懔υO(shè)定下，MEGA實(shí)現(xiàn)了121倍的推理加速及476倍的能效提升。
圖3. MEGA加速器架構(gòu)圖
更進(jìn)一步，研發(fā)團(tuán)隊(duì)提出了全流程圖神經(jīng)網(wǎng)絡(luò)訓(xùn)練加速系統(tǒng)FastGL，極大地提升了圖神經(jīng)網(wǎng)絡(luò)在超大規(guī)模圖數(shù)據(jù)上的訓(xùn)練速度。與最先進(jìn)的圖神經(jīng)網(wǎng)絡(luò)訓(xùn)練加速系統(tǒng)DGL對(duì)比，F(xiàn)astGL在多種超大規(guī)模圖數(shù)據(jù)和圖神經(jīng)網(wǎng)絡(luò)模型上實(shí)現(xiàn)了平均2.2倍的加速效果，并且具有更好的擴(kuò)展性，有利于應(yīng)對(duì)未來(lái)圖數(shù)據(jù)規(guī)模持續(xù)增加的趨勢(shì)。
圖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ò)高效計(jì)算方案，也為解決通用大模型推理中由超大規(guī)模參數(shù)量和超長(zhǎng)KV cache引起的計(jì)算、訪存雙密集加速等難題提供了啟發(fā)。
發(fā)表該系列研究的國(guó)際會(huì)議HPCA、ASPLOS與ISCA、MICRO并稱(chēng)為計(jì)算機(jī)體系結(jié)構(gòu)領(lǐng)域的四大頂會(huì)，是計(jì)算機(jī)體系結(jié)構(gòu)研究的風(fēng)向標(biāo)，ICLR則為深度學(xué)習(xí)領(lǐng)域的頂級(jí)學(xué)術(shù)會(huì)議。
論文信息：
[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)開(kāi)源代碼：
A2Q: https://github.com/weihai-98/A-2Q
FastGL: https://github.com/a1bc2def6g/fastgl-ae
