一、解決大圖內(nèi)存/計(jì)算問題的三個(gè)范式
在兩年前做的tutorial里面，我們有介紹過關(guān)于大規(guī)模神經(jīng)網(wǎng)絡(luò)，并且對20年以前的大規(guī)模圖神經(jīng)網(wǎng)絡(luò)的進(jìn)展有過一些介紹。在那個(gè)時(shí)候，考慮的是這樣三個(gè)范式：layer wise，node wise layer wise和graph wise sampling。
現(xiàn)在來看，歸根結(jié)底是要去減少圖數(shù)據(jù)在內(nèi)存和計(jì)算上的需求。最簡單的方法是對圖進(jìn)行采樣。回顧一下當(dāng)年的一些總結(jié)，從14年的圖神經(jīng)網(wǎng)絡(luò)開始走進(jìn)人們的視野，到17年GCN的爆火，其實(shí)一直以來，對于大規(guī)模圖神經(jīng)網(wǎng)絡(luò)的研究都是一個(gè)非常連續(xù)的過程。大家都是在朝著如何構(gòu)造更好的采樣和如何減少采樣造成的偏差兩個(gè)方向思考問題，也涌現(xiàn)出了非常多的優(yōu)秀工作。
我們真的解決了大規(guī)模GNN的問題嗎？我的答案是解決了，但沒有真正解決。首先，確實(shí)解決了在實(shí)際工業(yè)中的應(yīng)用，尤其是基于子圖采樣的方法，永遠(yuǎn)都可以采樣出一個(gè)子圖，Apply一個(gè)很復(fù)雜的模型，最后得到一個(gè)合適的預(yù)測。這個(gè)在騰訊的一些業(yè)務(wù)場景，比如推薦，已經(jīng)有了很好的實(shí)踐。
但是這個(gè)問題并沒有真正的解決，因?yàn)檫@個(gè)方法其實(shí)是回避了核心問題，不能真正在大圖做GNN更新。在真正做實(shí)踐的時(shí)候會發(fā)現(xiàn)，由于各個(gè)地方的系統(tǒng)可能不一樣，數(shù)據(jù)存儲格式不一樣，圖采樣的效率本身會依賴于系統(tǒng)實(shí)現(xiàn)。而圖采樣的時(shí)間消耗，很可能比訓(xùn)練的消耗更大。另外，這種采樣會帶來精度下降和信息缺失的風(fēng)險(xiǎn)。尤其是在制藥和生物的一些場景里面，是不能隨便的對比進(jìn)行采樣的。
那么近兩年大規(guī)模圖神經(jīng)網(wǎng)絡(luò)的進(jìn)展到底怎樣呢？可以總結(jié)為一句話，“我不想去做采樣，但是要把大規(guī)模的GNN給做了”。
二、針對大規(guī)模GNN的優(yōu)化
這里仿照之前WWW的GNN Tutorial做了這樣的一個(gè)圖，這個(gè)總結(jié)可能不是非常的全面，是我個(gè)人對這塊領(lǐng)域的總結(jié)。接下來就挑一些重要的點(diǎn)來進(jìn)行介紹，這兩年大家到底做了什么事情？簡單來說，首先我們用了圖系統(tǒng)里面的一些分布式圖系統(tǒng)計(jì)算的一些概念。把傳統(tǒng)的GAS范式進(jìn)行了擴(kuò)展，成了SAGA。基于這個(gè)范式，就會有很多需要系統(tǒng)優(yōu)化的點(diǎn)，那么具體優(yōu)化的點(diǎn)可能就存在于：第一，圖劃分和圖劃分的優(yōu)化；第二，對于節(jié)點(diǎn)特征傳輸?shù)膬?yōu)化；第三，對于流水線和通訊的優(yōu)化。接下來就每個(gè)單獨(dú)的進(jìn)行簡要的介紹。
1、對傳統(tǒng)圖計(jì)算模型GAS擴(kuò)展
首先，什么是SAGA，在說什么是SAGA的時(shí)候，我們就會討論什么是GAS。GSA是當(dāng)年圖計(jì)算里面一個(gè)分布式圖計(jì)算里面一個(gè)非常經(jīng)典的范式，它把整個(gè)圖計(jì)算劃分成了三個(gè)步驟：Gather，Apply和Scatter 。
- Gather是什么意思呢？Gather就是說去通過邊來收集鄰居的信息。
- Apply是什么意思呢？Apply就相當(dāng)于是把收集到的信息去計(jì)算出更新的這個(gè)節(jié)點(diǎn)信息。
- Scatter的意思就是把更新后的節(jié)點(diǎn)信息更新到這個(gè)邊上面。
很多圖算法都可以通過GAS的方式進(jìn)行Formalize，比如PageRank。也就是說面對基于消息傳播的圖神經(jīng)網(wǎng)絡(luò)的時(shí)候，就可以基于GAS方式進(jìn)行擴(kuò)展，叫做SAGA。
其中可以分為四個(gè)步驟：Scatter、Apply Edge、Gather和Apply Vertex。
Scatter 和原來的GAS的Scatter 是一樣的，就是把節(jié)點(diǎn)的數(shù)據(jù)先輸送到邊上面。這一個(gè)步驟是GPU Intensive。
然后再把這個(gè)節(jié)點(diǎn)的這個(gè)特征輸?shù)竭吷先ィ院罂赡軙數(shù)竭@個(gè)邊上面的消息進(jìn)行一些處理，這個(gè)處理有可能，比如說是GAT里面的計(jì)算權(quán)重，或者其他一些復(fù)雜操作，會Apply一個(gè)神經(jīng)網(wǎng)絡(luò)去處理，這個(gè)肯定是GPU intensive。
第三個(gè)， Gather傳過來消息了，并且可能是做了處理的消息，那就要通過鄰居的關(guān)系，把消息進(jìn)行匯聚，得到新的消息，那么這肯定就是CPU intensive。
第四個(gè)步驟，得到匯聚后的消息后，可能還需要一個(gè)額外的Update，可能是通過一個(gè)神經(jīng)網(wǎng)絡(luò)進(jìn)行，那么這個(gè)Update的結(jié)果其實(shí)就是最終的更新后的節(jié)點(diǎn)的表示。那么這個(gè)步驟叫做Apply Vertex，那么這其實(shí)就是GPU Intensive。
那么其實(shí)可以從這個(gè)范式里面可以看到：
第一，我們可以吸取以前GS的一些優(yōu)良的系統(tǒng)，優(yōu)化的一些策略，比如一些圖劃分、Pipeline、調(diào)度的策略。
第二，這里也挑戰(zhàn)，比如前文提到這是一個(gè)GPU和CPU交替intensive的一個(gè)任務(wù)，那么怎么優(yōu)化？比如說GPU和CPU之間的數(shù)據(jù)傳輸，甚至是不同集群之間的數(shù)據(jù)傳輸，其實(shí)是一個(gè)非常重大的問題，也是近年來研究的熱點(diǎn)。那么接下來介紹一下它們之間的關(guān)系。
2、圖劃分&圖劃分優(yōu)化
（1）單機(jī)上的圖劃分
首先，對于第一個(gè)優(yōu)化點(diǎn)，圖劃分和圖劃分的優(yōu)化，本質(zhì)上就是其實(shí)因?yàn)闆]辦法把圖放到顯存和單機(jī)上面，因此就需要把圖劃分到不同的Server上面去。對于Newer Graph來說，采用這種Locality Aware的圖劃分策略，相當(dāng)讓連在同一個(gè)節(jié)點(diǎn)上的邊能盡可能在同一個(gè)窗口，這樣再去做更新的時(shí)候，就不用訪問其他的窗口，就可以優(yōu)化內(nèi)存的訪問。
（2）基于模型圖劃分
第二個(gè)，ROC工作其實(shí)引入了Linear Regression Model，這個(gè)Linear Regression Model會去預(yù)測每一次Server里面的執(zhí)行時(shí)間，然后通過這個(gè)執(zhí)行時(shí)間去更新下一輪Partition的圖Partition的策略。這里是基于模型的Cost、Running Time的圖的劃分。
（3）縱向劃分節(jié)點(diǎn)特征
P3這篇干脆就回避了劃分的一些缺點(diǎn)，因?yàn)槿绻腔诠?jié)點(diǎn)或邊的分割，會導(dǎo)致額外的信息通訊和信息損失，所以P3的劃分就不是基于圖的，而是基于特征的。這個(gè)motivation的特征維度特別大，沒辦法全部放到同一個(gè)單機(jī)上面，但這個(gè)節(jié)點(diǎn)本質(zhì)上是一個(gè)Adjacency List，是可以放到每一個(gè)單獨(dú)的Server上面去的，因此可以對于每一個(gè)Server節(jié)點(diǎn)的特征進(jìn)行動縱向劃分。Per dimension把這些劃分放到各個(gè)機(jī)器上面去，這樣既保證了圖的完整性，也保證了Somehow減少信息通訊的代價(jià)。
3、對于節(jié)點(diǎn)特征傳輸?shù)膬?yōu)化
很多工作都利用了節(jié)點(diǎn)特征傳輸?shù)膬?yōu)化，會發(fā)現(xiàn)有些時(shí)候節(jié)點(diǎn)的原始特征是比GNN的Dimension大很多，尤其是當(dāng)節(jié)點(diǎn)是一個(gè)C的特征，通過之后的Embedding，這個(gè)節(jié)點(diǎn)特征會非常非常大。所以如果基于節(jié)點(diǎn)特征在節(jié)點(diǎn)的機(jī)器和機(jī)器之間做通信是非常不利的，不經(jīng)濟(jì)的。怎么盡量去減少這種Move節(jié)點(diǎn)的這樣的通信，或者說計(jì)算，也就是是對于節(jié)點(diǎn)特征傳輸?shù)囊粋€(gè)優(yōu)化？那么這個(gè)地方也是有三項(xiàng)工作。
PA graph其實(shí)就是一個(gè)靜態(tài)緩存的機(jī)制。比如我們知道某個(gè)顯存可能只有20個(gè)G做計(jì)算，那顯存如果還有10個(gè)G，就會去緩存一些節(jié)點(diǎn)的原始特征。這個(gè)策略，本質(zhì)上是基于靜態(tài)的緩存機(jī)制，它會去把邊的點(diǎn)的Out Degree進(jìn)行排序，會把Out Degree多的點(diǎn)扔到緩存里面去，那顯然Out Degree多的點(diǎn)很容易有更高的概率參與計(jì)算，所以說開始會把節(jié)點(diǎn)特征傳輸?shù)耐ㄐ帕繙p少。
在DistGNN里面，設(shè)計(jì)了一個(gè)更復(fù)雜的開始Blocking的機(jī)制。這個(gè)開始Blocking的機(jī)制會把節(jié)點(diǎn)分為兩類——要更新的目標(biāo)節(jié)點(diǎn)和目標(biāo)節(jié)點(diǎn)鄰居節(jié)點(diǎn)。傳統(tǒng)的方法是基于目標(biāo)節(jié)點(diǎn)進(jìn)行遍歷，然后去拉它的鄰居節(jié)點(diǎn)DistGNN。這里說的機(jī)制則是反向去遍歷鄰居節(jié)點(diǎn)，然后對目標(biāo)節(jié)點(diǎn)特征進(jìn)行alternately的更新。這樣做好處是目標(biāo)節(jié)點(diǎn)的特征會放到CPU的緩存里面，每次去找鄰居節(jié)點(diǎn)，只需要讀一次，不需要緩存，每次去更新目標(biāo)節(jié)點(diǎn)的時(shí)候，都是在CPU緩存里面更新。這樣子就是可以減少反復(fù)拉取鄰居節(jié)點(diǎn)的開銷。
對于P3來說，這個(gè)設(shè)計(jì)就會更復(fù)雜一些，因?yàn)镻3的目標(biāo)是不希望產(chǎn)生任何原始特征之間拉取的通信，所以設(shè)計(jì)了一個(gè)hybrid parameter的機(jī)制。對于輸入層的GNN來說，模型會并行在每一個(gè)單機(jī)，也就是每一個(gè)單機(jī)有一部分特征。計(jì)算出sub partial activation，然后每個(gè)機(jī)器在本地計(jì)算完partial activation，會把這個(gè)partial activation匯聚成一個(gè)輸入層的GNN的輸出，得到這個(gè)輸入層的GNN輸出以后，再走正常的Data Parallelism，去做這樣計(jì)算的動機(jī)是原始特征的維度特別大，那么第一層GNN的參數(shù)和通信量都會非常大。而除了第一層以外，其他GNN的黑洞Size其實(shí)都會比較小，這種時(shí)候如果在本地算好了第一層再去后面去算更深層的DNN，通信代價(jià)能大為減少。
4、流水線、通信優(yōu)化
這是一個(gè)非常巧妙的設(shè)計(jì)，對于這種流水線通信的優(yōu)化，就更不用說了。流水線通信優(yōu)化，其實(shí)也是傳統(tǒng)的圖計(jì)算里面做的非常多的，簡單來說，同步更新的一些信息需要算完一輪特征以后再更新下一輪的特征，但有些時(shí)候可以不這樣去做，而是一個(gè)異步更新，并且在異步更新的過程中，使數(shù)據(jù)聚合的過程中，或者在更新的過程中使用一些老的特征，同時(shí)也會設(shè)計(jì)一個(gè)老化機(jī)制。
這樣相當(dāng)于是在這個(gè)分布式系統(tǒng)里面的半同步的機(jī)制，通過半同步機(jī)制，可以很好的去設(shè)計(jì)這個(gè)流水線，并且減少多機(jī)通信的代價(jià)。這里因?yàn)闀r(shí)間關(guān)系，就不具體展開講了。
要注意的是，可以用的老的特征，但不能用太老的特征，如果這個(gè)特征太老，也會暫停更新，等其他節(jié)點(diǎn)的特征更新到最近的epoch以后，才繼續(xù)更新。
其實(shí)像Dorylus這種工作，本身就是在多個(gè)機(jī)器上面，而且每個(gè)機(jī)器的這個(gè)什么角色還不一樣，分為Graph Server，Lambda thread和Permit Server，這種情況下，對Pipeline流水線的設(shè)計(jì)是有很多細(xì)節(jié)的。總的來說，現(xiàn)有的方法，也是慢慢從單機(jī)多GPU到多機(jī)混合的趨勢，從經(jīng)驗(yàn)上的靜態(tài)劃分到這種動態(tài)的劃分，以及引入更多系統(tǒng)層面的Pipeline的優(yōu)化。
同時(shí)，像SANCUS的工作，其實(shí)也還進(jìn)一步去證明了異步的更新是可以保證模型收斂和通訊優(yōu)化的。就是說慢慢的大家從一些簡單的實(shí)踐到復(fù)雜實(shí)踐，從一些沒有理論保證的實(shí)驗(yàn)，到一些理論保證實(shí)踐。這一塊的發(fā)展還是非常的迅速的。
四、未來方向
接下來談一下對未來方向的一個(gè)理解。首先我們發(fā)現(xiàn)，實(shí)際上SAGA的這個(gè)范式并不適用于所有圖神經(jīng)網(wǎng)絡(luò)。如果不是基于Message Passage神經(jīng)網(wǎng)絡(luò)，其實(shí)它就不符合SAGA范式，同樣的，就沒辦法利用到已有的GAS范式里面系統(tǒng)優(yōu)化的一些trick，來對系統(tǒng)的整個(gè)代價(jià)進(jìn)行優(yōu)化，比如說對于Graph Transformer這樣的模型，最近也非常的火，從20年開始到現(xiàn)在也有很多這樣的模型出現(xiàn)。
那么，能不能在這樣的模型里做Full Graph的優(yōu)化、或者說Full Graph的訓(xùn)練，其實(shí)是一個(gè)非常有挑戰(zhàn)性的問題。而且這并不是說沒有應(yīng)用，比如對于蛋白質(zhì)建模來說，實(shí)際上如果蛋白質(zhì)的這個(gè)序列足夠長，把這個(gè)蛋白質(zhì)作為一個(gè)Graph的話，它的訓(xùn)練代價(jià)會非常的大。顯存可能會出現(xiàn)暴漲，那么怎么在這種非Message Passing的框架下面去對這種大規(guī)模GNN做系統(tǒng)優(yōu)化，是一個(gè)非常重要的課題。
第二個(gè)就是很多時(shí)候，以前的這個(gè)圖神經(jīng)網(wǎng)絡(luò)里面是沒有包含幾何信息的，什么叫幾何信息，就是說節(jié)點(diǎn)可能是有空間位置信息的，比如坐標(biāo)、速度，或者說一些其他的幾何信息，這些東西實(shí)在很多，尤其是AI for Science領(lǐng)域是非常常見的。
比如我們對于粒子進(jìn)行建模。對于這種催化劑系統(tǒng)的模擬，都會包含幾何信息，而幾何信息本身的查詢和更新就是數(shù)據(jù)庫領(lǐng)域的一個(gè)非常重要的課題，這個(gè)和Spatial Database有很強(qiáng)的聯(lián)系，那么對于這種類型的數(shù)據(jù)，我們在已有的等變神經(jīng)網(wǎng)絡(luò)成果上面，能不能對它進(jìn)行系統(tǒng)化或者規(guī)模化？因?yàn)閷?shí)際上這個(gè)系統(tǒng)化，規(guī)模化也是非常急需的，就比如之前做的一個(gè)催化劑的比賽OCP，這個(gè)比賽里面都包含幾百萬個(gè)催化劑系統(tǒng)，數(shù)據(jù)量是T級別的，實(shí)際上對于模型的訓(xùn)練和推理都有非常大的挑戰(zhàn)，那么對于這種幾何信息的輸出神經(jīng)網(wǎng)絡(luò)，是不是有很好的解決的方法，這也是未來研究的方向。
五、Q&A
Q：目前要把圖神經(jīng)網(wǎng)絡(luò)從學(xué)術(shù)界的東西變成工業(yè)界真正能使用的東西，還有很大的差距。其中不僅是算法本身的優(yōu)化，還有現(xiàn)在些創(chuàng)業(yè)企業(yè)都在做graph，做computing platform，然后做一個(gè)真正的end to end。大家想要從系統(tǒng)優(yōu)化到上層算法，再到application，完整的做一套系統(tǒng)出來，然后來服務(wù)drug discovery、金融等領(lǐng)域。比如用這種更先進(jìn)方法來從廣告品中挖掘到更多信息。您對算法架構(gòu)整個(gè)這塊都有很多研究了，您對graph plus，也就是graph platform computing as a service，怎么看？
A：我覺得這是非常有前景的一個(gè)方向，簡單來說，實(shí)際上graph computing本身的platform，在AI時(shí)代之前就有很多人研究了。因?yàn)樵趃raph上面有很多傳統(tǒng)的圖算法，也是需要去做分布式的計(jì)算，這種分布式的計(jì)算就是service的，但是以前我們做的可能都是一些基礎(chǔ)的計(jì)算。而現(xiàn)在像圖神經(jīng)網(wǎng)絡(luò)等技術(shù)興起后，我們發(fā)現(xiàn)這個(gè)系統(tǒng)里面會面臨更多的挑戰(zhàn)。
比如說以前的數(shù)據(jù)只需要在CPU上面算就可以了。現(xiàn)在類似于GNN這樣的結(jié)構(gòu)，我們一定會涉及到一種混合架構(gòu)，第一做起來很難；第二，這個(gè)東西如果做出來了，門檻很高，所以說這一塊是非常有前景的一個(gè)方向。
而且這一塊市場的玩家目前不是很多，尤其是對于大企業(yè)來說，這一塊要去集合做AI的和做系統(tǒng)的人，這兩塊人湊在一起，在大企業(yè)里面其實(shí)是比較困難的，因?yàn)楫吘勾笃髽I(yè)都是以商業(yè)為主軸的。所以我們正是需要一些創(chuàng)業(yè)的公司把這個(gè)東西做出來以后，真正去服務(wù)一些實(shí)際的應(yīng)用。比如你說的drug discovery，或者說一些AI + Science，或者說可能的一些在城市道路，甚至社交網(wǎng)絡(luò)上面的應(yīng)用。其實(shí)這塊東西屬于門檻高，有前景，處于還需要人進(jìn)一步去開拓和做這樣的一個(gè)狀態(tài)。
