一、解決大圖內(nèi)存/計算問題的三個范式

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

現(xiàn)在來看，歸根結(jié)底是要去減少圖數(shù)據(jù)在內(nèi)存和計算上的需求。最簡單的方法是對圖進行采樣?；仡櫼幌庐斈甑囊恍┛偨Y(jié)，從14年的圖神經(jīng)網(wǎng)絡開始走進人們的視野，到17年GCN的爆火，其實一直以來，對于大規(guī)模圖神經(jīng)網(wǎng)絡的研究都是一個非常連續(xù)的過程。大家都是在朝著如何構(gòu)造更好的采樣和如何減少采樣造成的偏差兩個方向思考問題，也涌現(xiàn)出了非常多的優(yōu)秀工作。

(資料圖片僅供參考)

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

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

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

二、針對大規(guī)模GNN的優(yōu)化

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

1、對傳統(tǒng)圖計算模型GAS擴展

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

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

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

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

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

然后再把這個節(jié)點的這個特征輸?shù)竭吷先?，以后可能會對輸?shù)竭@個邊上面的消息進行一些處理，這個處理有可能，比如說是GAT里面的計算權(quán)重，或者其他一些復雜操作，會Apply一個神經(jīng)網(wǎng)絡去處理，這個肯定是GPU intensive。

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

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

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

第一，我們可以吸取以前GS的一些優(yōu)良的系統(tǒng)，優(yōu)化的一些策略，比如一些圖劃分、Pipeline、調(diào)度的策略。

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

2、圖劃分&圖劃分優(yōu)化（1）單機上的圖劃分

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

（2）基于模型圖劃分

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

（3）縱向劃分節(jié)點特征

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

3、對于節(jié)點特征傳輸?shù)膬?yōu)化

很多工作都利用了節(jié)點特征傳輸?shù)膬?yōu)化，會發(fā)現(xiàn)有些時候節(jié)點的原始特征是比GNN的Dimension大很多，尤其是當節(jié)點是一個C的特征，通過之后的Embedding，這個節(jié)點特征會非常非常大。所以如果基于節(jié)點特征在節(jié)點的機器和機器之間做通信是非常不利的，不經(jīng)濟的。怎么盡量去減少這種Move節(jié)點的這樣的通信，或者說計算，也就是是對于節(jié)點特征傳輸?shù)囊粋€優(yōu)化？那么這個地方也是有三項工作。

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

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

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

4、流水線、通信優(yōu)化

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

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

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

其實像Dorylus這種工作，本身就是在多個機器上面，而且每個機器的這個什么角色還不一樣，分為Graph Server，Lambda thread和Permit Server，這種情況下，對Pipeline流水線的設計是有很多細節(jié)的?？偟膩碚f，現(xiàn)有的方法，也是慢慢從單機多GPU到多機混合的趨勢，從經(jīng)驗上的靜態(tài)劃分到這種動態(tài)的劃分，以及引入更多系統(tǒng)層面的Pipeline的優(yōu)化。

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

四、未來方向

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

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

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

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

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

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

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

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

關(guān)鍵詞：