從信用卡交易、社交網絡到推薦系統,再到生物學中的運輸網絡和蛋白質間的相互作用,圖形是建模和分析復雜連接的首選數據結構。圖形神經網絡 (GNN) 具有學習和推理圖形結構化數據的能力,已在各個領域成為改變游戲規則的技術。
但是,發現這些圖形中隱藏的模式和寶貴見解可能具有挑戰性,尤其是在 GNN 的數據采樣和端到端訓練中。
為了彌補這一差距,NVIDIA 推出了GNN 框架,這是專為 DGL 和 PyG 設計的容器,具備以下特點:
- 用于數據采樣的 GPU 加速
- GNN 訓練和部署工具 (GNN 工具)
本文概述了 NVIDIA 加速 DGL 和 PyG 容器的優勢,展示了客戶在生產環境中使用這些容器的方式,并提供了性能指標。
在 NGC 目錄中介紹 DGL 容器
深度圖形庫 (DGL) 是一款熱門的開源庫,它可以在現有的深度學習框架(如 PyTorch)上實現和訓練圖神經網絡 (GNN)。
我們很高興地宣布,DGL 現已通過其他 NVIDIA 庫進行加速,并作為 容器?通過 NGC 目錄 提供 —— 這是一個 GPU 加速的 AI/ML、HPC 應用程序、SDK 和工具的中心。該目錄使用戶能夠更快地訪問性能優化的軟件,并簡化構建和部署 AI 解決方案的過程,幫助您更快地將產品推向市場。有關更多信息,請參閱 NGC 目錄中用于 AI、數字孿生和 HPC 的數百個預訓練模型(視頻)。
23.09 版本的 DGL 容器提高了#DGL 用戶的數據采樣和訓練性能。以下是此版本的主要功能。
用于數據加載程序采樣的 GPU 加速
RAPIDS cuGraph 的 sampler 能夠在幾秒鐘內處理數千億個邊,并同時為數千批世界上最大的 GNN 數據集計算樣本。DGL 容器隨附 cuGraph-DGL,這是 DGL 的加速擴展,使用戶能夠利用這一驚人的性能。
根據使用 8 個 V100 GPU 運行的基準測試,即使在中型數據集(約 10 億個邊緣)上,cuGraph 數據加載性能也至少比原生 DGL 快 2 – 3 倍。cuGraph-DGL 采樣通過在多個節點和多個 GPU 中分配圖形,還能在此過程中節省內存,比線性擴展高達 1000 億個邊緣更好。
cuGraph 可以在 16 秒內對 1000 億個邊緣進行采樣! NVIDIA 專有庫 cuGraph-ops 加速了 cuGraphSAGE、cuGraphGAT 和 cuGraphRGCN 等 GNN 運算符和模型,將模型前向時間縮短了一半。
GNN 訓練和部署工具
GNN 工具是一個靈活的平臺,可輕松訓練和部署 GNN 模型。此工具基于熱門的 Deep Graph Library (DGL) 和 PyTorch Geometric (PyG) 框架構建,使您能夠為快速 GNN 實驗構建端到端工作流。
它提供了一個完全模塊化和可配置的工作流,可為自定義 GNN 用例實現快速迭代和實驗。 NVIDIA 在我們的容器中包含示例 notebook,以便輕松進行實驗。
多架構支持
NGC 中發布的 DGL 容器具有 x86 和 ARM64 版本,以支持新的 NVIDIA Grace Hopper GPU 這兩個版本使用相同的容器標簽。從基于 Arm 的 Linux 系統中提取容器時,您需要提取 ARM64 容器。
使用基于 ARM64 的 DGL 容器在 Grace Hopper 上訓練 GNN 的示例
GNN 訓練中的統一虛擬尋址 (UVA) 模式極大地受益于 NVIDIA Grace CPU 和 Grace Hopper 中 NVIDIA Hopper GPU 之間的連接。在 Grace Hopper 上,使用 ogn-papers100M 數據集訓練同一 GraphSAGE 模型需要 1.9 秒/次,與使用 PCIe 連接的 H100*Intel CPU 訓練相比,速度快了約 9 倍(表 1)。
系統 | GH200 | H100+Intel CPU | A100+AMD (DGX A100) |
AMD Genoa (僅限 CPU) |
(秒/次) | 1.9 | 16.92 | 24.8 | 107.11 |
這些數字利用 Grace Hopper 上的巨大頁面來處理圖形及其功能。它們在 512 GB Grace Hopper 節點上進行基準測試。該模型運行時批量大小為 4096,扇出 (3030) (在兩層 GraphSAGE 模型中查看每個節點多達 30 個相鄰節點)。這在 CUDA 12.1 上的 DGL 版本 1.1 上運行。
DGL 訓練性能
訓練 GNN 面臨的挑戰之一是數據加載過程。在某些情況下,例如使用 GraphSAGE 對 ogb-papers100M 數據集進行節點分類,數據加載過程占端到端訓練時間的 90%以上。DGL 0.8 v 啟用 UVA 模式,以實現圖形特征的高效 GPU 加載,這提高了性能。
考慮使用 GraphSAGE 模型進行 ogbn-products 節點分類任務的數據集。該數據集包含約 240 萬個節點和 6190 萬條邊。在 DGX-1 V100 GPU 上,采用 UVA 模式進行訓練時,相比僅使用 CPU,它可以提供高達 20 倍的速度提升(見圖 1)。

由于可以將 ogbn-product 數據集加載到 GPU 顯存中,速度甚至可以加快 115 倍(圖 2)。

對于大型數據集(例如 ogn-papers100M),必須開啟 UVA 模式。
圖 3 顯示了節點分類任務中 ogn-papers100M 的每輪訓練時間(以秒為單位)。它具有 1.11 億個節點和 32 億個邊緣。

您可以在 NGC DGL 23.09 容器的/workspace/examples/
多 GPU 目錄。
PyTorch Geometric 容器
PyTorch Geometric (PyG) 是另一個熱門的開源庫,用于編寫和訓練各種應用的 GNN。我們將推出使用 cuGraph 等 NVIDIA 庫加速的 PyG 容器。
根據使用 8 個 A100 GPU 運行的基準測試,在中型數據集(大約 10 億個邊緣)上,cuGraph 數據加載性能至少比原生 PyG 快 4 倍。
我們已經觀察到一些客戶從 PyG 容器中受益,并且我們計劃利用 PyG 加速來與 NVIDIA BioNeMo 模型一起使用。
客戶成功案例
以下介紹了不同的公司如何使用 NVIDIA 加速 DGL 和 PyG 容器來加速其工作流程。
適用于基于物理的 ML 的 GNN
NVIDIA Modulus 是一個開源框架,用于在 Python 中構建、訓練和微調基于物理性質的 機器學習(ML) 模型。
隨著 GNN 在計算流體動力學、分子動力學模擬和材料科學等學科中的應用興趣與日俱增, NVIDIA Modulus 開始利用 DGL 和 cuGraphOps 庫來支持 GNN.
NVIDIA Modulus 目前支持包括 MeshGraphNet、AeroGraphNet 和 GraphCast 在內的 GNN,用于基于網格的模擬和全球天氣預報的模型。除了網絡架構,Modulus 還提供了一系列工具,用于開發 天氣預報、空氣動力學模擬 以及 渦流脫落 的示例。
用于欺詐檢測的 GNN
憑借數十年的經驗,美國運通 在使用 AI 驅動的工具和模型監控及降低欺詐風險方面有著顯著的成就。它們還能有效地識別信用卡行業中參與欺詐活動的個人。
American Express AI Labs 正在持續開展研究,通過實施基于圖形的機器學習解決方案,更深入地了解欺詐者網絡。
NGC 上發布的 DGL 容器使 AmEx 能夠試驗各種 GNN 架構,并大規模利用節點和邊緣信息。這使得 AmEx 能夠在多節點、多 GPU 環境中使用 NVIDIA 庫處理數百萬個節點和數十億個邊緣時實現計算效率。此外,用戶友好且適應性強的庫使他們能夠輕松自定義組件,例如損失函數、采樣技術等。
用于藥物研發的 GNN
Astellas 作為一家以前瞻性藥物研發方法而聞名的領先制藥公司,正在利用 GNN 的功能來完成藥物研發中的一系列關鍵任務。這些任務包括使用基于擴散模型的概率性深度學習模型來生成 3D 分子構象、進行特征提取和構建預測性 ML 模型,特別是在從頭開始的蛋白質設計和工程方面。
因此,Astellas 利用 NVIDIA PyG 和 DGL 容器的計算效率,在藥物研發活動中促進和放大基于 GNN 的 AI/ML 解決方案。由 PyG 和 DGL 提供支持的 AI 和基于 ML 的流程的集成使 Astellas 能夠提高其內部能力,為研究人員和 AI 從業者提供開發和實施 AI 驅動的切割技術所需的專業知識。
作為一個用例,與基于仿真的傳統 3D 分子構象方法相比,Astellas 科學家可以實現至少 50 倍的加速率。
Genentech,作為一家生物技術公司以及 Roche Group 的一部分,正在結合使用 GNN 和 NVIDIA PyG 容器,以加速其小分子預測訓練。PyG 容器為 PyG 框架提供了穩定的基礎,使得開發團隊能夠更加專注于開發工作,而不是搭建開發環境。
后續步驟
借助 NVIDIA 加速的 DGL 和 PyG 容器,您還可以顯著提高 GNN 的數據采樣和訓練性能。
要開始使用,請下載以下資源:
?