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

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

圖1. 圖數據具有復雜的拓撲結構

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

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

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

為此，在專(zhuān)用推理架構方面，團隊針對壓縮算法特性和圖數據特性設計了一款專(zhuān)用加速器MEGA。通過(guò)比特位自適應的存儲方式和高效的在線(xiàn)編解碼壓縮方法，最大限度減少了存儲碎片化，降低了數據傳輸開(kāi)銷(xiāo)。此外，MEGA采用的Condense-Edge調度執行方式，能夠將稀疏的邊連接聚合在一起以增強訪(fǎng)問(wèn)節點(diǎn)特征時(shí)的空間局部性，提升數據復用度。在與英偉達 RTX3090 GPU相當的峰值算力設定下，MEGA實(shí)現了121倍的推理加速及476倍的能效提升。

圖3. MEGA加速器架構圖

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

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

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

發(fā)表該系列研究的國際會(huì )議HPCA、ASPLOS與ISCA、MICRO并稱(chēng)為計算機體系結構領(lǐng)域的四大頂會(huì )，是計算機體系結構研究的風(fēng)向標，ICLR則為深度學(xué)習領(lǐng)域的頂級學(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

來(lái)源：中國科學(xué)院自動(dòng)化研究所