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