一、解決大圖內存/計算問題的三個范式

在兩年前做的tutorial里面，我們有介紹過關于大規模神經網絡，并且對20年以前的大規模圖神經網絡的進展有過一些介紹。在那個時候，考慮的是這樣三個范式：layer wise，node wise layer wise和graph wise sampling。

現在來看，歸根結底是要去減少圖數據在內存和計算上的需求。最簡單的方法是對圖進行采樣。回顧一下當年的一些總結，從14年的圖神經網絡開始走進人們的視野，到17年GCN的爆火，其實一直以來，對于大規模圖神經網絡的研究都是一個非常連續的過程。大家都是在朝著如何構造更好的采樣和如何減少采樣造成的偏差兩個方向思考問題，也涌現出了非常多的優秀工作。

我們真的解決了大規模GNN的問題嗎？我的答案是解決了，但沒有真正解決。首先，確實解決了在實際工業中的應用，尤其是基于子圖采樣的方法，永遠都可以采樣出一個子圖，Apply一個很復雜的模型，最后得到一個合適的預測。這個在騰訊的一些業務場景，比如推薦，已經有了很好的實踐。

但是這個問題并沒有真正的解決，因為這個方法其實是回避了核心問題，不能真正在大圖做GNN更新。在真正做實踐的時候會發現，由于各個地方的系統可能不一樣，數據存儲格式不一樣，圖采樣的效率本身會依賴于系統實現。而圖采樣的時間消耗，很可能比訓練的消耗更大。另外，這種采樣會帶來精度下降和信息缺失的風險。尤其是在制藥和生物的一些場景里面，是不能隨便的對比進行采樣的。

那么近兩年大規模圖神經網絡的進展到底怎樣呢？可以總結為一句話，“我不想去做采樣，但是要把大規模的GNN給做了”。

二、針對大規模GNN的優化

這里仿照之前WWW的GNN Tutorial做了這樣的一個圖，這個總結可能不是非常的全面，是我個人對這塊領域的總結。接下來就挑一些重要的點來進行介紹，這兩年大家到底做了什么事情？簡單來說，首先我們用了圖系統里面的一些分布式圖系統計算的一些概念。把傳統的GAS范式進行了擴展，成了SAGA。基于這個范式，就會有很多需要系統優化的點，那么具體優化的點可能就存在于：第一，圖劃分和圖劃分的優化；第二，對于節點特征傳輸的優化；第三，對于流水線和通訊的優化。接下來就每個單獨的進行簡要的介紹。

1、對傳統圖計算模型GAS擴展

首先，什么是SAGA，在說什么是SAGA的時候，我們就會討論什么是GAS。GSA是當年圖計算里面一個分布式圖計算里面一個非常經典的范式，它把整個圖計算劃分成了三個步驟：Gather，Apply和Scatter 。

• Gather是什么意思呢？Gather就是說去通過邊來收集鄰居的信息。
• Apply是什么意思呢？Apply就相當于是把收集到的信息去計算出更新的這個節點信息。
• Scatter的意思就是把更新后的節點信息更新到這個邊上面。

很多圖算法都可以通過GAS的方式進行Formalize，比如PageRank。也就是說面對基于消息傳播的圖神經網絡的時候，就可以基于GAS方式進行擴展，叫做SAGA。

其中可以分為四個步驟：Scatter、Apply Edge、Gather和Apply Vertex。

Scatter 和原來的GAS的Scatter 是一樣的，就是把節點的數據先輸送到邊上面。這一個步驟是GPU Intensive。

然后再把這個節點的這個特征輸到邊上去，以后可能會對輸到這個邊上面的消息進行一些處理，這個處理有可能，比如說是GAT里面的計算權重，或者其他一些復雜操作，會Apply一個神經網絡去處理，這個肯定是GPU intensive。

第三個， Gather傳過來消息了，并且可能是做了處理的消息，那就要通過鄰居的關系，把消息進行匯聚，得到新的消息，那么這肯定就是CPU intensive。

第四個步驟，得到匯聚后的消息后，可能還需要一個額外的Update，可能是通過一個神經網絡進行，那么這個Update的結果其實就是最終的更新后的節點的表示。那么這個步驟叫做Apply Vertex，那么這其實就是GPU Intensive。

那么其實可以從這個范式里面可以看到：

第一，我們可以吸取以前GS的一些優良的系統，優化的一些策略，比如一些圖劃分、Pipeline、調度的策略。

第二，這里也挑戰，比如前文提到這是一個GPU和CPU交替intensive的一個任務，那么怎么優化？比如說GPU和CPU之間的數據傳輸，甚至是不同集群之間的數據傳輸，其實是一個非常重大的問題，也是近年來研究的熱點。那么接下來介紹一下它們之間的關系。

2、圖劃分&圖劃分優化

（1）單機上的圖劃分

首先，對于第一個優化點，圖劃分和圖劃分的優化，本質上就是其實因為沒辦法把圖放到顯存和單機上面，因此就需要把圖劃分到不同的Server上面去。對于Newer Graph來說，采用這種Locality Aware的圖劃分策略，相當讓連在同一個節點上的邊能盡可能在同一個窗口，這樣再去做更新的時候，就不用訪問其他的窗口，就可以優化內存的訪問。

（2）基于模型圖劃分

第二個，ROC工作其實引入了Linear Regression Model，這個Linear Regression Model會去預測每一次Server里面的執行時間，然后通過這個執行時間去更新下一輪Partition的圖Partition的策略。這里是基于模型的Cost、Running Time的圖的劃分。

（3）縱向劃分節點特征

P3這篇干脆就回避了劃分的一些缺點，因為如果是基于節點或邊的分割，會導致額外的信息通訊和信息損失，所以P3的劃分就不是基于圖的，而是基于特征的。這個motivation的特征維度特別大，沒辦法全部放到同一個單機上面，但這個節點本質上是一個Adjacency List，是可以放到每一個單獨的Server上面去的，因此可以對于每一個Server節點的特征進行動縱向劃分。Per dimension把這些劃分放到各個機器上面去，這樣既保證了圖的完整性，也保證了Somehow減少信息通訊的代價。

3、對于節點特征傳輸的優化

很多工作都利用了節點特征傳輸的優化，會發現有些時候節點的原始特征是比GNN的Dimension大很多，尤其是當節點是一個C的特征，通過之后的Embedding，這個節點特征會非常非常大。所以如果基于節點特征在節點的機器和機器之間做通信是非常不利的，不經濟的。怎么盡量去減少這種Move節點的這樣的通信，或者說計算，也就是是對于節點特征傳輸的一個優化？那么這個地方也是有三項工作。

PA graph其實就是一個靜態緩存的機制。比如我們知道某個顯存可能只有20個G做計算，那顯存如果還有10個G，就會去緩存一些節點的原始特征。這個策略，本質上是基于靜態的緩存機制，它會去把邊的點的Out Degree進行排序，會把Out Degree多的點扔到緩存里面去，那顯然Out Degree多的點很容易有更高的概率參與計算，所以說開始會把節點特征傳輸的通信量減少。

在DistGNN里面，設計了一個更復雜的開始Blocking的機制。這個開始Blocking的機制會把節點分為兩類——要更新的目標節點和目標節點鄰居節點。傳統的方法是基于目標節點進行遍歷，然后去拉它的鄰居節點DistGNN。這里說的機制則是反向去遍歷鄰居節點，然后對目標節點特征進行alternately的更新。這樣做好處是目標節點的特征會放到CPU的緩存里面，每次去找鄰居節點，只需要讀一次，不需要緩存，每次去更新目標節點的時候，都是在CPU緩存里面更新。這樣子就是可以減少反復拉取鄰居節點的開銷。

對于P3來說，這個設計就會更復雜一些，因為P3的目標是不希望產生任何原始特征之間拉取的通信，所以設計了一個hybrid parameter的機制。對于輸入層的GNN來說，模型會并行在每一個單機，也就是每一個單機有一部分特征。計算出sub partial activation，然后每個機器在本地計算完partial activation，會把這個partial activation匯聚成一個輸入層的GNN的輸出，得到這個輸入層的GNN輸出以后，再走正常的Data Parallelism，去做這樣計算的動機是原始特征的維度特別大，那么第一層GNN的參數和通信量都會非常大。而除了第一層以外，其他GNN的黑洞Size其實都會比較小，這種時候如果在本地算好了第一層再去后面去算更深層的DNN，通信代價能大為減少。

4、流水線、通信優化

這是一個非常巧妙的設計，對于這種流水線通信的優化，就更不用說了。流水線通信優化，其實也是傳統的圖計算里面做的非常多的，簡單來說，同步更新的一些信息需要算完一輪特征以后再更新下一輪的特征，但有些時候可以不這樣去做，而是一個異步更新，并且在異步更新的過程中，使數據聚合的過程中，或者在更新的過程中使用一些老的特征，同時也會設計一個老化機制。

這樣相當于是在這個分布式系統里面的半同步的機制，通過半同步機制，可以很好的去設計這個流水線，并且減少多機通信的代價。這里因為時間關系，就不具體展開講了。

要注意的是，可以用的老的特征，但不能用太老的特征，如果這個特征太老，也會暫停更新，等其他節點的特征更新到最近的epoch以后，才繼續更新。

其實像Dorylus這種工作，本身就是在多個機器上面，而且每個機器的這個什么角色還不一樣，分為Graph Server，Lambda thread和Permit Server，這種情況下，對Pipeline流水線的設計是有很多細節的。總的來說，現有的方法，也是慢慢從單機多GPU到多機混合的趨勢，從經驗上的靜態劃分到這種動態的劃分，以及引入更多系統層面的Pipeline的優化。

同時，像SANCUS的工作，其實也還進一步去證明了異步的更新是可以保證模型收斂和通訊優化的。就是說慢慢的大家從一些簡單的實踐到復雜實踐，從一些沒有理論保證的實驗，到一些理論保證實踐。這一塊的發展還是非常的迅速的。

四、未來方向

接下來談一下對未來方向的一個理解。首先我們發現，實際上SAGA的這個范式并不適用于所有圖神經網絡。如果不是基于Message Passage神經網絡，其實它就不符合SAGA范式，同樣的，就沒辦法利用到已有的GAS范式里面系統優化的一些trick，來對系統的整個代價進行優化，比如說對于Graph Transformer這樣的模型，最近也非常的火，從20年開始到現在也有很多這樣的模型出現。

那么，能不能在這樣的模型里做Full Graph的優化、或者說Full Graph的訓練，其實是一個非常有挑戰性的問題。而且這并不是說沒有應用，比如對于蛋白質建模來說，實際上如果蛋白質的這個序列足夠長，把這個蛋白質作為一個Graph的話，它的訓練代價會非常的大。顯存可能會出現暴漲，那么怎么在這種非Message Passing的框架下面去對這種大規模GNN做系統優化，是一個非常重要的課題。

第二個就是很多時候，以前的這個圖神經網絡里面是沒有包含幾何信息的，什么叫幾何信息，就是說節點可能是有空間位置信息的，比如坐標、速度，或者說一些其他的幾何信息，這些東西實在很多，尤其是AI for Science領域是非常常見的。

比如我們對于粒子進行建模。對于這種催化劑系統的模擬，都會包含幾何信息，而幾何信息本身的查詢和更新就是數據庫領域的一個非常重要的課題，這個和Spatial Database有很強的聯系，那么對于這種類型的數據，我們在已有的等變神經網絡成果上面，能不能對它進行系統化或者規模化？因為實際上這個系統化，規模化也是非常急需的，就比如之前做的一個催化劑的比賽OCP，這個比賽里面都包含幾百萬個催化劑系統，數據量是T級別的，實際上對于模型的訓練和推理都有非常大的挑戰，那么對于這種幾何信息的輸出神經網絡，是不是有很好的解決的方法，這也是未來研究的方向。

五、Q&A

Q：目前要把圖神經網絡從學術界的東西變成工業界真正能使用的東西，還有很大的差距。其中不僅是算法本身的優化，還有現在些創業企業都在做graph，做computing platform，然后做一個真正的end to end。大家想要從系統優化到上層算法，再到application，完整的做一套系統出來，然后來服務drug discovery、金融等領域。比如用這種更先進方法來從廣告品中挖掘到更多信息。您對算法架構整個這塊都有很多研究了，您對graph plus，也就是graph platform computing as a service，怎么看？

A：我覺得這是非常有前景的一個方向，簡單來說，實際上graph computing本身的platform，在AI時代之前就有很多人研究了。因為在graph上面有很多傳統的圖算法，也是需要去做分布式的計算，這種分布式的計算就是service的，但是以前我們做的可能都是一些基礎的計算。而現在像圖神經網絡等技術興起后，我們發現這個系統里面會面臨更多的挑戰。

比如說以前的數據只需要在CPU上面算就可以了。現在類似于GNN這樣的結構，我們一定會涉及到一種混合架構，第一做起來很難；第二，這個東西如果做出來了，門檻很高，所以說這一塊是非常有前景的一個方向。

而且這一塊市場的玩家目前不是很多，尤其是對于大企業來說，這一塊要去集合做AI的和做系統的人，這兩塊人湊在一起，在大企業里面其實是比較困難的，因為畢竟大企業都是以商業為主軸的。所以我們正是需要一些創業的公司把這個東西做出來以后，真正去服務一些實際的應用。比如你說的drug discovery，或者說一些AI + Science，或者說可能的一些在城市道路，甚至社交網絡上面的應用。其實這塊東西屬于門檻高，有前景，處于還需要人進一步去開拓和做這樣的一個狀態。