導讀 本文主要介紹飛槳圖學習的大模型訓練框架。
介紹會圍繞下面三點展開：
1. 圖學習框架發展趨勢
2. 圖學習框架的 GPU 加速
3. 工業級大規模圖表示學習 PGLBox
分享嘉賓｜黃正傑 百度 資深研發工程師
編輯整理｜王超 智聯招聘
出品社群｜DataFun
01
首先圖是描述複雜世界的一個通用語言。包括我們在網際網路裡面的一些像社交網路，它是人和人之間的連線。以及近期一些比較火熱的像生物計算的一些領域，它是一些化學分子的圖的結構。以及像知識圖譜，它是一些實體與關係的連線。推薦系統是使用者與物品之間的一些關聯。
圖神經網路這個概念也是在近年來比較火的。之前在建模圖的時候，通常會基於一個譜域的概念，就是一個 Spectral-based 的方法。會首先把一些圖透過傅立葉變換轉化到頻域，再去頻域上操作。當時可能是因為這個理念比較複雜，並且它對於大圖的擴充套件性比較差，所以關注這方面工作的人比較少。大概在 2017 年左右，像圖卷積網路，還有一些像 Message Passing 的網路出現之後，大家會轉向用更簡單的架構去看待圖網路，例如 Spatial-based 的方法，就是在空域上去理解圖神經網路。我們可以類比到二維的圖卷積裡面去，如果你要建模一個點的元素，它是跟你周圍的畫素有關的。在圖神經網路裡面的話就是如果我要建模一個點、它的一個屬性以及表示，通常只會跟它的鄰居有關，然後就透過一些鄰居的計算，我們就可以得到我們中心節點的一個表示。
在基於空域的圖神經網路裡面，我們通常是關注兩個點，一個是我們怎麼做鄰居的聚合去得到中心點的表示；另一個是我們有一張大圖了，怎麼去透過節點的匯聚得到整個圖結構的表達。
目前主流的圖神經網路基本都可以描述成這種叫做訊息傳遞的形式，即 Message Passing。像 DGL 還有 PYG 這些目前比較熱門的圖神經網路框架，包括我們的 PGL 也是沿用這樣基於訊息傳遞的正規化去定義圖神經網路。這裡的 Message Passing 就是說我們只需要建模鄰居是怎麼把的訊息從源節點傳送到一個目標節點，以及目標節點獲取到源節點的訊息之後怎麼做，例如做一個加權的求和或者是做一個均值。透過這樣的一個框架，我們目前基本上能實現大部分的圖神經網路，也甚至包括現在的 Transformer，假設把 Transformer 理解成一個全連線的圖神經網路，其實 Transformer 的結構也是可以透過這種訊息傳遞形式去實現。
上圖是我們的一個圖學習的框架 PGL。大概是在兩年前釋出的，主要是基於飛槳的 PaddlePaddle 深度學習框架，在上面搭建了一些包括圖的引擎，以及圖網路的一些程式設計的介面，上面我們也放了很多預置的模型，供大家去使用。感興趣的同學可以去我們的Github，可以 pip install 裝我們的框架，以及跑一些示例。
圖神經網路大規模的訓練，目前是有一套比較成熟的標準流程。特別是大圖的情況下，我們通常會沿用 2017 年的一篇論文叫做 GraphSAGE，它主要的理念就是訓練在大圖裡面的某些帶標籤的節點，我們通常會對某個點的鄰居進行取樣 sample，後面用 aggregation 就是不停地聚合或者透過訊息傳遞的圖神經網路去代表我們中心的一個點。對於不同的框架，我們都會主要解決三個問題，第一個是我們的圖是怎麼存的，第二個是我們的子圖是怎麼取樣，包括一些特徵是怎麼從圖裡面拉取出來。最後是高層圖網路的一些訊息傳遞演算法是怎麼定義的。
相比於一些大模型，在圖網路裡面大部分的消耗反而是在圖的儲存以及特徵的拉取上面。比起剛剛說到的一些大模型引數量上，目前模型在純網路的模組是會比較小的，但即使在這種情況下，我們的圖神經網路訓練的 GPU 利用率還是很低的。最大的瓶頸就在於 CPU 上圖的取樣的瓶頸，以及我們從 CPU 到 GPU 的資料傳輸的瓶頸。因為我們的圖很大，所以通常情況下 GPU 是存不下來的。我們會把圖和一些取樣的演算法實現在記憶體，最後取樣完之後，我們再把它挪到我們的 GPU 去做計算。
隨著硬體的逐漸地發展，我們也會發現一個趨勢，就是現在的視訊記憶體越來越大，包括是 8 塊的 A100 的 80G 我們能在單臺機器上組到一個 640g 的視訊記憶體。那麼這是一個什麼樣的概念？我舉一個例子，就是現有的最大的一個異構圖的資料集，Open Graph Benchmark 裡面最大的一張圖是叫 MAG240M，裡面是一些論文作者引用網路，大概是 2 億的點 10 億的邊。它裡面最大的東西是 310G 的特徵，就是使用 BERT 抽取的文字特徵，每個文字擁有 768 維，這裡有大概 300 多 G 特徵。但是在這麼大的情況下，我們的 A100 80g，如果我們把 8 個卡拼在一起，已經能超過這麼一個資料集的容量了。那麼我們是不是可以做這樣的一個事情？就是 GPU 儲存這麼大，我們是不是可以把我們整個圖演算法的流程完全地 GPU 化滿足大部分的場景。因此我們也在想怎麼對我們圖學習的框架進行 GPU 的加速。
說到加速的話，因為大圖很難實現。那我們先把概念轉到圖很小的情況下，我們會怎麼做？在圖很小的情況下，一個很直接的辦法就是我們把圖的結構完全的塞到視訊記憶體裡面去。再透過在視訊記憶體裡面實現的一些圖取樣的運算元，那麼就可以把我們在圖學習演算法裡面的圖儲存以及取樣的資料流全部的挪到我們的 GPU 上面去，就可以達到很極致的速度。但是這裡最大的問題就是實際上在做的時候，我們的圖還是挺大的。如果是大概 200G 的一個圖結構，一個單卡的 GPU 已經做不到了。
一個折中的辦法就是我們能不能用 GPU 取樣的運算元訪問記憶體的圖資料，做到大圖的擴充套件。這個工作是 2021 年，大概六七月份的時候有個叫 Torch-Quiver 的團隊他們做了一個事情，就是把記憶體當做視訊記憶體的一塊，用一個叫做 UVA 的模式，用 GPU 的取樣運算元，直接對記憶體訪問去做取樣。達到的效果就是會比這純 CPU 的取樣快很多，比 GPU 視訊記憶體的儲存要慢一些，但是圖就已經可以變得很大了。
他們具體的做法就是把 GPU 的圖結構存到鎖頁記憶體裡，透過一種叫做統一地址的 Unified Virtual Addressing 技術，把記憶體當視訊記憶體，使用 GPU 取樣的 Kernel 就可以直接訪問統一的一個地址。他們給出了一組資料，就是當用 UVA 的情況下做取樣，會相比純 CPU 的接近快 20 倍的速度。 
除了 UVA 模式下單卡用記憶體配合，那麼如果我們在多卡的情況下，剛剛也說到了我們 8 塊 A100 是能達到 640G 的。我們也是在想能不能用到例如 GPU 之間的通訊NVLink 技術，用 GPU 間通訊傳輸大的頻寬，頂替掉記憶體到視訊記憶體上面的傳輸。
這裡我們針對圖裡面最大的一塊特徵按列切割的形式進行分卡的儲存。例如我們左邊的這裡有一個 300G 的圖特徵的儲存，我們把它切到每一張卡，大概每一張卡是 40G 的視訊記憶體消耗。那麼我們就可以把一整個 MAG240 加上圖結構在 8 個 A100 上面儲存，組成了一個很大的共享的 Tensor。
在實際的訓練裡面，我這裡用兩個 GPU 來舉例子，當我們在做資料並行的時候，每個 GPU 只是負責某一部分節點的訓練，取樣它子圖 ID 的時候，我們可以按列用一次的 All-gather，把所有的點對應的一塊子圖，全部透過 All-gather 獲取，然後在區域性的按列切分的特徵上面把一些 feature 獲取。獲取之後再透過 All2all 去交換特徵，把整個圖的結構特徵還原。透過這樣的 All2all 把整個特徵切分，用 GPU 做傳輸，我們目前可以把 R-GAT 在 MAG240 上面半分鐘就可以訓完一個 epoch 了，相比之前的相對於 CPU 的或者說 DGL 的那種也是 CPU 模式的速度會更快。
在 2021 年 KDD 比賽裡面我們提了一個方法叫 R-UniMP，裡面做了很多對於異構圖取樣的最佳化，以及對於圖裡面一些標籤應用的最佳化。當時我們的模型做得特別複雜，當時我們是訓練 100 個 epoch 得到 SoTA 成績是大概要 40 個小時。今年我們透過一系列的最佳化，在 8 卡的 A100 上 1.1 小時就能出一個去年最好的模型，也給我們今年參加 NEURIPS22 的比賽提供了很大的演算法調研的空間。剛剛說到的幾個場景都是偏向學術界的一些怎麼去訓練大圖的場景。通常來講，學術界可能關注的是節點分類，節點分類的資料結構是很整齊的，也不需要大規模的節點特徵。
在工業界的場景，特別是網際網路的產品應用，我們在做推薦的時候，通常會面臨一個事情，就是對我們的使用者以及物品去建模。推薦裡面有兩個常用的演算法，一個是 Item-based 協同過濾，就是去衡量兩個物品之間的一些相似，以及 User-base 協同過濾，就是去衡量兩個使用者之間的相似度。
通常圖學習在網際網路裡面的一些應用大部分都是基於這樣的情況，就是我們會在海量的日誌裡面使用者以及物品構造一張很大的圖，包括有社會的關係、使用者的一些行為以及物品的一些關聯。我們透過這些圖表示學習的一些方法對每個點給一個向量，在這個向量空間裡面可以用距離去度量，就是不同的物品以及使用者之間的一些相似度。
在傳統的一些演算法上面，包括像矩陣分解，也是可以做這樣的向量表示的，以及像幾年前 Deepwalk 和 Meta-Path 這樣的演算法，就是透過隨機遊走加 Word2vec 的演算法，也可以構造出圖的表示。今年來比較火熱的就是用圖對比學習的方法，透過圖的資料增強，透過對比學習對兩個子圖去抽取使用者和物品的行為表示。作為一個圖學習框架，最重要的是需要滿足我們業務調研的一些需求，我們也整理了目前大部分的圖學習的方法，發現基本上是可以抽象為下面四個步驟的。
第一個是大家選用的圖可能是不同的，例如 item 的一種 session graph 就是同構圖來的，使用者物品可能是個二部圖。有些像異構圖的話，就可能是在同一個圖裡面有點選關係，喜好關係或者一些購買關係，那它是一種複雜的異構圖的形式。在有了圖之後，做一些圖上面的自監督度或無監督任務的話，我們通常會用一些隨機遊走的方法來產生一些訓練的樣本。有了樣本之後，對於每一個點我們會採用不同的形式做表達，例如我們會做不同鄰居的取樣，包括像一個比較出名的演算法叫 PinSage，就是透過隨機遊走選鄰居，不同的演算法會有自己的定義。在資料構成之後，最後才是我們的圖網路的結構，就是我們怎麼去選不同的結構去建模鄰居和自身的表達。目前基本上所有的無監督的圖學習的方法，都可以抽象成這麼四個流程。
2021年，在我們 PGL 內部，我們提供了一套叫 Graph4Rec 的一個工具，可以透過五項的配置，就是配置好圖資料，配置圖的隨機遊走樣本，取樣的演算法，正負樣本取樣生成的方法，以及網路選擇，就能自動生成一套表示學習的訓練，得到每個節點高質量的 embedding。但是之前我們這套方法是基於分散式的 CPU 去訓練的。
在大規模應用場景下，就出現了現在這樣的一些痛點，例如我們通常會有一些千億級的稀疏特徵，包括我們用來存一些使用者的 ID、圖文 ID 以及一些離散化的圖文影片的屬性。在這種情況下，我們通常需要一套大規模引數伺服器，在訓練的時候要跟引數伺服器進行大量 embedding 的一些查詢及互動。2021年我們也還是用 CPU 引數伺服器加 CPU 圖引擎去配合一個 CPU 的 MPI叢集做訓練。它最大的問題在於現階段不同模態的建模越來越相似，包括用 Transformer 的一些結構的話，是無法去滿足一些複雜的模型的。
其次就是當我們的圖取樣鄰居特別多的時候，拉取 embedding 以及圖的取樣的情況下，它通訊量很大。我們也經常會發現，例如我們搞 20 臺機器去訓練的時候，只要有一臺通訊被打掛了，這個訓練就很容易失敗。有人也可能會說就是能不能用 CPU 引數伺服器加圖引擎，去配合 GPU 去訓練，滿足剛剛說到的一些複雜模型的需求。但是我們也會發現 IO 是很密集的，就是一些引數伺服器的通訊，也是會導致整體的 GPU 利用率很低。5. PGLBox：全流程 GPU 化的大規模圖學習引擎
延續我們最開始說到的，就是既然現在的 GPU 伺服器已經例如 8 塊 A100 我們有 640G 的一個視訊記憶體，甚至記憶體可能都已經有三個 T 的情況下，那麼我們是不是就可以做一套全流程完全 GPU 化的大規模圖學習引擎？就是把我們的圖結構、圖的特徵表以及我們的引數伺服器全部塞到一臺機器的 GPU 裡面去，即透過 GPU 的分卡，全部塞到我們的 GPU 裡面去。那麼就把我們剛剛說到的所有的跨機的通訊全部取消掉。在這個情況下，640G 已經能支援很大的一個量級的業務了。所以這裡我們也是實現了一個全流程 GPU 最佳化的一個大規模圖學習引擎，也是為了方便我們的一些業務。因為跟學術界的不一樣，就是大家的資料很規整，我們的業務可能每次過來的就是使用者給兩個 ID，我們就要去構圖，或就做查詢。所以我們也做了易用性上面的很多個 GPU 的雜湊表，去支援雜湊的構圖以及雜湊的圖查詢，也透過 GPU 加速這種複雜異構圖的取樣。以及我們在 GPU 上面做了很多的引數伺服器的梯度最佳化策略，來減少視訊記憶體的使用。6. PGLBox：透過層次化儲存和流水線，支援百億節點/數百億邊
如果業務再大一點，就是我們的圖可能再大到百億點或者百億邊以上的規模，我們也提供了另外一種層次化的策略，就是把我們的整一個大圖存到一個 SSD 以及記憶體上面去，當我們使用到一個訓練的時候，再把對應的東西挪到視訊記憶體裡面去。這裡的圖表示學習訓練我們做了一系列取樣的情況下，通常我們訓練的一塊它是很集中的，就集中在某一塊的子圖。所以在這裡我們就引入了一個叫 Pass 的一個概念，就是我們一個 epoch 可能要把整個圖遊走，全部的點給遍歷了一遍。但是我們的一個 Pass 就是把整一個 epoch 把它拆成若干個小的 Pass，在 Pass 裡面做樣本的遊走，以及引數伺服器的一些特徵的準備。準備完了之後，這一塊就很適合直接在視訊記憶體裡面做。那麼在視訊記憶體裡面我們就一直訓了 N 個 Batch，把這整個 Pass 訓練完之後，我們再把對應的一些 embedding 以及吸收的一些引數同步到一個 SSD 以及記憶體的引數伺服器上面去。不同的 Pass 之間，我們可以做一個流水線的並行。就例如我們在做 pull embedding，從記憶體去拉取對應的一個 embedding 的時候，第二個 Pass 的遊走的取樣以及樣本的生成和分析已經可以做了。所以透過這樣的一個 Pass 級的一個訓練並行，儘管我們引入了這些層次化的儲存，速度也是跟全視訊記憶體的模式能做到一個持平的效果。
有了這麼一個 PGLBox 的大圖的表示學習的框架。我們在公司內也做了很多業務的升級。包括像原來有些 Deepwalk 模型，可能是 480 分鐘能做完的，現在已經可以一個小時內就解決了。更復雜的模型，像 GraphSAGE 這種的就是會隨著我們取樣的鄰居個數，導致計算量指數上漲的，在子圖結構的指數上漲的同時，特徵的拉取以及通訊量也是在指數上升的。在 GPU 這種單機 8 卡的情況下，就基本上能把這些通訊給最佳化掉，像這種複雜的 GraphSAGE 一樣的圖網路，也能去做到 28 倍的提速。除此之外也是順應時代潮流，我們透過這種全域性 GPU 化能做到一些複雜模型的支援，像我們應用場景裡面可能有不同的一些複雜關係，我們的節點本身是多模態的，裡面可能有使用者、文字和影像，大規模的引數伺服器是提供了一些長期行為 ID 的建模，預訓練模型就提供了這些跨模態內容的理解。因為有了這種 GPU 計算的圖架構，我們就可以把以往預訓練的大模型去跟這些大規模引數伺服器放在同一套框架下面訓起來，做一個端對端的學習，包括同時去調節預訓練大模型裡面的一些引數以及大規模的基於 ID 的一些引數伺服器。透過目前這套 PGLBox 的框架，我們實現了單機 8 卡百億級規模圖表示學習訓練，相比 CPU 分散式方案提升 28 倍。
PGLBox 目前已經在github開源，具體地址如下：
問答環節
Q1：GPU 雜湊表跟普通的 CPU 雜湊表有什麼區別？A1：目前 GPU 雜湊表用起來可能沒有 CPU 的那麼成熟，所以目前那些主流的開源框架，像 DGL 還有 PYG 那些都沒有做這個事情。我們做這個雜湊表，它跟 CPU 的也沒有太大的區別，主要就是功能上要有這麼一個功能，用來簡化在 GPU 上做圖。Q2：Cache 方式是提前挪一些到 GPU 上嗎？還是全部放到 GPU上？A2：Cache 方式我們是提前挪一些到 GPU 上。像 MAC240 這種異構圖的資料，它全部加起來也就 300 多 G，是因為它每一維可能有 768 維的特徵。那如果沒有那麼大，例如我們平時做表示學習的 100 維，甚至在做比較早期的推薦系統的一些場景，很多人可能做 8 維，這裡可以放很大的一張圖了。在這種情況下，圖可以做得很大。所以我們基本上就卡到一個程度，一個 Pass 剛好能把視訊記憶體完全佔滿，留一部分的空間去做模型的矩陣計算，基本就可以了。Q3：Pass 之間的流水線並行會導致梯度的過期嗎？A3：目前流水線的話，梯度上面我們會就每一次從 CPU 拉回來的東西是序列的，保證 Pass 去拉 embedding 的時候，下一個 Pass 去拉 embedding 是必須在我這個 Pass 的 push 之後的，所以這裡是沒有問題的。具體能並行掉的是遊走取樣的那個樣本，就是樣本的構成，以及取樣遊走樣本，這一塊也挺佔時間的。Q4：多機的擴充套件的問題
A4：單機 8 卡的話在卡內就可以用 NVLink 這些技術去做。多機的話還涉及到 embedding，例如跨機裡面會包含一些網路頻寬的一些問題。
來自 “ ITPUB部落格 ” ，連結：http://blog.itpub.net/70027824/viewspace-2944662/，如需轉載，請註明出處，否則將追究法律責任。