導讀
近期阿里雲機器學習平臺 PAI 團隊和達摩院 GraphScope 團隊聯合推出了面向 PyTorch 的 GPU 加速分佈式 GNN 框架 GraphLearn-for-PyTorch(GLT) 。GLT 利用 GPU 的強大並行計算性能來加速圖採樣，並利用 UVA 來減少頂點和邊特徵的轉換和拷貝。對於大規模圖，GLT 使用了生產者-消費者的架構，通過異步併發的分佈式採樣和特徵查找以及熱點緩存功能支持在多個 GPU 或多個機器上進行高效的分佈式訓練。接口上，GLT 保持了 PyTorch的風格，並且和 PyG 兼容，只需少量代碼修改就可以加速 PyG 的單機訓練程序，或者將 PyG 單機模型改成分佈式訓練。此外，GLT 還提供了靈活的分佈式訓練部署以滿足不同的需求。
開源地址：https://github.com/alibaba/graphlearn-for-pytorch
文檔地址：https://graphlearn-torch.readthedocs.io/en/latest/index.html
背景介紹
圖神經網絡作爲一種圖數據上表示學習的方法已經被廣泛應用在圖相關的各個領域，在電商推薦、安全風控、生物分子等領域取得了實際落地。圖神經網絡由於其獨特的數據處理邏輯和神經網絡計算邏輯，需要有專門的學習框架來支持訓練。PAI團隊之前開源了大規模工業級分佈式圖學習框架 GraphLearn（https://github.com/alibaba/graph-learn）。GraphLearn 以 TensorFlow 1.x 系列爲主，採用 ps 架構的異步訓練模式，支持十億節點，百億邊規模的大規模異構圖分佈式訓練，應用於阿里內外部若干業務場景。隨着PyTorch 的流行，其更加靈活的貼近 Python 的接口，簡單易調試等特性使得算法開發者更傾向於使用 PyTorch 開發模型。DGL 和 PyG等基於PyTorch的開源GNN框架以單機爲主，無法支持大規模圖的分佈式訓練。
此外，由於 GPU 並行計算的優勢，圖神經網絡使用 GPU 訓練比 CPU 訓練有數倍的提升。然而常見的圖神經網絡框架將圖拓撲數據和特徵存在內存裏，使用CPU進行採樣和特徵查找並將數據拷貝到GPU進行神經網絡訓練，這個過程中圖採樣和特徵查找部分很容易成爲整體訓練的瓶頸。下面我們以大規模圖上典型的訓練流程爲例對訓練過程的性能瓶頸進行分析說明。
一個典型的GNN訓練流程[1] 包括:
- 子圖拓撲採樣，採樣多跳鄰居並組成子圖；
- 查詢子圖裏節點或者邊的特徵；
- 將子圖格式轉換成神經網絡訓練需要的格式並且拷貝到GPU顯存中；
- 對原始特徵進行處理，比如離散特徵進行embedding lookup；
- 鄰居聚合；
- 節點更新。
其中，3和4爲可選步驟。常見的 GNN 模型神經網絡參數相對來說比較小，因此計算量也比較小，瓶頸通常在1-4步，具體來說主要是I/O操作，包括通信，數據拷貝和格式轉換。這導致即使使用GPU進行訓練，GPU的資源利用率也很低，整體吞吐以及擴展性很難提高。
綜上所述，我們需要一個高效的基於PyTorch 的分佈式 GNN 訓練框架，能夠充分利用 GPU 等硬件資源，能夠基於圖的數據分佈性質，結合不同算法模型和並行策略做相應優化，減少通信和數據轉換耗時，提升整體吞吐。
關鍵設計
設計初期，我們和 Quiver[2] 團隊合作針對 GPU 採樣的可行性進行了初步探索，發現 GPU 採樣相比 CPU 能夠帶來數量級的提升。此外，爲了減少特徵在 CPU 內存和 GPU 顯存之間的拷貝開銷，特徵可以直接存儲在 GPU 顯存裏。然而對於規模比較大的數據來說，顯存無法存儲所有特徵數據，因此我們用 UVA 技術來進行特徵的存儲，相比直接內存存儲也有數量級的性能提升。工業界的圖規模很容易突破單機的極限，因此我們進一步設計了分佈式訓練框架。具體來說，我們使用生產者-消費者範式來處理分佈式圖採樣和特徵查找以及模型訓練的關係，將圖採樣、特徵查找與訓練進行解耦，使用多進程並行和協程異步併發來加速採樣和特徵查找，並使用預取和熱點緩存的方式進一步減少訓練端的等待，提升端到端吞吐。考慮到用戶遷移成本和易用性，在接口上我們保持了和 PyG 的兼容，只需少量改動 PyG 代碼就可以加速 PyG 的訓練，或者將其遷移到 GLT 的分佈式訓練上。以下爲我們具體闡述幾個關鍵的設計點。
GPU採樣
GLT 將圖拓撲使用 CSR 格式存儲在 GPU 顯存或者 pin memory裏，實現了 CUDA 採樣算子來進行 GPU 並行採樣。使用 CSR 存儲可以很容易得到每個節點的鄰居，並且獨立對每個節點進行採樣，因此可以方便地利用 GPU 多線程進行並行採樣。我們使用了蓄水池算法來進行無放回隨機採樣。在batch size大的情況下，GPU 比 CPU 採樣能有數量級的提升。
UnifiedTensor
爲了消除CPU 內存到 GPU 顯存的拷貝開銷，一個比較直觀的方法是將特徵存放在 GPU 顯存裏，然而由於單卡的顯存有限，在特徵數據比較大的情況下也很難完全把特徵存到顯存裏。因此，在GLT中我們利用圖自身的特性如power law分佈和採樣訪問特性如有些度比較高的節點被訪問的概率高，將部分熱點特徵存放在 GPU 顯存裏，其他特徵存在內存，同時需要利用 UVA 讓 GPU 訪問內存裏的特徵。GLT 設計了 UnifiedTensor 將 CUDA Tensor 和 CPU Tensor 統一管理起來，以提供簡潔高效的數據訪問。進一步，如果 GPU 之間可以直接進行 peer2peer 訪問（具有NVLink），這些 GPU 的顯存也可以被統一管理，從而擴大特徵在顯存的存儲。GLT 使用UnifiedTensor 將這些不同硬件設備上的存儲統一管理起來，提供直接訪問 CUDA Tensor，通過 NVLink 訪問其他 GPU 上的 CUDA Tensor，並通過 UVA 進行 ZERO-COPY 訪問 CPU Tensor的能力，上層查找元素接口就像普通 Tensor 一樣，底層會自動去對應的設備上進行訪存操作。
Feature
Feature 由 UnifiedTensor 構成，具有硬件拓撲感知功能。具體來說，首先，按照用戶指定 CPU/GPU 內存大小，對特徵進行劃分，分爲 GPU(hot) 部分和 CPU(cold) 部分。其次，對 GPU 部分，根據用戶指定的 replica 策略，進行 replica，包括每個卡 replica 和每個 NVLink 連接的 GPU group 之間的 replica。GPU group replica 的方式，相比卡間 replica, 可以有更多的 hot data 存在 GPU 上，因爲 GPU group 裏 GPU 之間都是可以 p2p 訪問的。實現上GLT抽象出 DeviceGroup 來統一表示卡間 replica 和組間 replica。一個 DeviceGroup 表示一組 NVLink 連接的 GPUs。假設8卡沒有 NVLink，那麼會對應8個 DeviceGroup，如果 GPU 0-3 兩兩 NVLink 連接，GPU 4-7 兩兩 NVLink 連接，那麼 GPU 0-3 爲一組 DeviceGroup, GPU 4-7 爲一組 DeviceGroup。實際測試中，使用UnifiedTensor 的Feature性能比 CPU Tensor的查找(包括拷貝到GPU) 快1個數量級，而且可以通過控制 GPU 存儲部分的比例來靈活達到速度和顯存佔用的平衡。
分佈式設計
GLT分佈式GNN訓練主要分成：分佈式採樣，特徵查找， 模型訓練3部分。一次採樣的結果一般比較小(最大爲十幾MB)，特徵查找的結果比較大(百MB)，訓練時使用特徵查找的結果進行神經網絡計算。對特徵查找來說需要考慮減少和訓練任務之間的數據轉換和拷貝。採樣和訓練之間是典型的生產者和消費者關係，因此可以分成不同任務，通過緩衝區連接，平衡生產者和消費者的處理能力，起到一個數據緩存的作用，同時也達到了一個解耦的作用。基於生產者-消費者方式，GLT的分佈式訓練有兩種基本類型的進程：採樣進程和訓練進程。
採樣進程：負責分佈式鄰居採樣和特徵收集。採樣結果將被髮送到採樣消息通道，該通道將進一步用於訓練任務。
訓練進程：對應於PyTorch DDP的分佈式訓練進程，通常，每個訓練進程將佔用一個GPU進行訓練。
這些進程可以靈活地分佈在不同的機器上，爲了更好地管理分佈式進程部署，GLT的分佈式訓練提供了兩種參考部署模式：Worker 模式和 Server-Client 模式。
Worker模式裏，數據切分後，每個機器持有一個分片，採樣進程和訓練進程一起部署在這些機器上。每個訓練進程可以spawn出多個採樣子進程，採樣進程通過一個共享內存的消息通道將採樣結果傳遞給訓練進程。對於採樣進程來說，可以使用多進程進行採樣，並且每個分佈式採樣算子都使用Python協程來併發執行，將結果放到消息通道里。爲了減少消息通道到訓練進程 GPU 的拷貝耗時，消息通道也可以放到pin memory上。
Server-Client模式下，集羣中存在兩種類型的機器節點，即 Server 節點和 Client 節點。採樣進程部署在 Server 節點，訓練進程分佈在所有 Client 節點上。採樣進程生成的樣本結果將通過一個RPC實現的遠程消息通道發送到當前訓練進程進行訓練。Server-Client 模式可以將採樣和訓練不同 workload 的任務放到不同機器，進行資源上的解耦。
總體架構
GLT 的主要目標是充分利用 GPU/NVLink/RDMA 等硬件資源和 GNN 模型的特性，加速單機和分佈式環境下的端到端 GNN 訓練。
存儲：在 GPU 訓練場景，圖採樣和 CPU-GPU 數據傳輸很容易成爲主要性能瓶頸。爲了加速圖採樣和特徵查找，GLT 實現了 UnifiedTensor 統一 CPU 和 GPU 的內存管理。爲了減少特徵收集引起的 CPU-GPU 數據傳輸開銷，GLT 支持將熱點特徵緩存在 GPU 顯存中，並通過 UVA 訪問其餘特徵數據。我們還利用高速 NVLink 在GPU 之間擴展 GPU 緩存的容量。
圖操作算子：存儲之上，GLT 實現了包括鄰居採樣、負採樣、特徵查找、子圖採樣等同時支持 CPU 和 GPU 圖操作算子。
分佈式算子：對於分佈式訓練，爲防止遠程數據訪問阻塞模型訓練進程，GLT 在 PyTorch RPC 之上封裝了一個高效的 RPC 框架，並採用多進程並行和異步併發的圖採樣和特徵查找操作來隱藏網絡延遲並提高端到端訓練吞吐量。
接口：爲了降低 PyG 用戶的學習成本，GLT 的上層 API，如Sampler, Dataset和Loader，接口上都與PyG兼容。因此，PyG 用戶只需修改很少的代碼即可充分利用 GLT 的加速能力。
模型：由於 GLT 與 PyG 兼容，你可以使用幾乎任何 PyG 的模型作爲基礎模型，此外我們也提供了豐富的分佈式訓練示例。
系統性能
我們在一臺配備A100 GPU的機器進行單機擴展性測試，測試環境爲 CUDA 11.4、PyTorch 1.12 和 GLT 0.2.0rc2，下圖展示了鄰居採樣和特徵查找的總吞吐量。可以看出 GLT 有線性的擴展性(由於有NVLink，多卡的緩存容量更多，因此會存在超線性加速)。
此外，我們還測試了多機的分佈式採樣和特徵查找的擴展性。下圖展示了每個機器配備2個A100 GPU的環境下，2個機器和4個機器相比單個機器的吞吐量加速比。測試使用 CUDA11.4、PyTorch 1.12和 GLT 0.2.0rc2 進行。可以看出，2機到4機也有近線性的擴展性。
最後，我們測試了分佈式e2e的性能。我們在2機每機2卡A100的設置下和DGL做了初步對比(DGL版本0.9.1，GLT版本0.2.0).
結語
本文介紹了基於PyTorch的GPU加速分佈式GNN框架GraphLearn-for-PyTorch(GLT)，GLT提供了分佈式 GPU 訓練的優化加速能力，能夠充分利用 GPU 等硬件資源進行圖採樣和特徵查找等操作，具有線性擴展性。上層接口上和 PyG 兼容，可以很容易地加速 PyG 已有模型或者將已有模型改成分佈式版本。GLT 已經開源並且在PyG, GraphScope中都有示例，後面我們會持續開發優化，歡迎使用和貢獻！
[1] P3: Distributed Deep Graph Learning at Scale
[2] Quiver: Supporting GPUs for Low-Latency, High-Throughput GNN Serving with Workload Awareness
作者：艾寶樂
點擊立即免費試用雲產品 開啓雲上實踐之旅！
本文爲阿里雲原創內容，未經允許不得轉載。