在非CUDA硬件上運行幾何學習:基于Intel Gaudi-v2 HPU的PyTorch框架移植實踐
論文標題:PyTorch-based Geometric Learning with Non-CUDA Processing Units: Experiences from Intel Gaudi-v2 HPUs
arXiv:2507.01031 (cross-list from cs.LG)
PyTorch-based Geometric Learning with Non-CUDA Processing Units: Experiences from Intel Gaudi-v2 HPUs
Fanchen Bu, Kijung Shin
Comments: Conference paper: Accepted in Korea Computer Congress (KCC) 2025. The library is available at this https URL
Subjects: Machine Learning (cs.LG); Software Engineering (cs.SE)
研究背景:當圖學習遇上硬件新選擇
想象一下,我們生活在一個由"關系"構成的世界里:社交網絡中人與人的關注關系、化學分子中原子的連接方式、推薦系統中用戶與商品的互動記錄,這些都可以用"圖"這種數據結構來表示。而幾何學習,尤其是圖神經網絡(GNN),就像是一位擅長解讀這些關系的"翻譯官",能從圖數據中提取有價值的信息,在節點分類、鏈接預測等任務中大顯身手。
過去,Nvidia的CUDA顯卡就像幾何學習領域的"標準配置電腦",幾乎所有算法都為它優化。但隨著技術發展,我們有了新的"硬件選擇"——比如Intel的Gaudi HPU,它像一臺性價比更高的"新型電腦",在性能和能耗上表現亮眼。然而麻煩來了:原來為CUDA"定制"的軟件(如PyTorch幾何學習框架),在這臺"新電腦"上無法直接運行,需要我們當"系統工程師",解決兼容性問題。
創新點:為非CUDA硬件搭建"橋梁"
論文的獨特之處在于,它沒有僅僅停留在理論分析,而是動手打造了一套"跨平臺工具包":
-
核心操作重構:針對Gaudi HPU不支持的散射、稀疏索引等基礎操作,用PyTorch重新實現,就像為新硬件編寫"驅動程序"。
-
全流程實踐指南:提供16個手把手教程和11個真實案例,從環境搭建到模型訓練,覆蓋節點分類、點云分析等場景,相當于一份"新電腦使用說明書"。
-
問題診斷數據庫:記錄移植中遇到的錯誤(如RuntimeError)和解決方法,比如通過移除torch.compile解決兼容性問題,形成"故障排除手冊"。
研究方法和思路:三步拆解移植工程
第一步:識別"不兼容零件"
- 發現torch-scatter(散射操作)、torch-sparse(稀疏矩陣)、torch-cluster(圖聚類)等庫不支持Gaudi HPU,就像發現新電腦無法識別舊軟件的安裝包。
第二步:打造"兼容適配器"
- 散射操作:用純PyTorch函數重寫,比如通過索引聚合實現"散射相加",讓新硬件能執行基礎數據聚合。
- 稀疏操作:將稀疏矩陣乘法轉換為散射操作,就像把"方言指令"翻譯成新硬件能聽懂的"普通話"。
- 高級圖操作:先在CPU上計算k-NN搜索等任務,再把結果傳給Gaudi HPU,類似"分工合作"的工作模式。
第三步:編寫"使用教程"
- 基于PyTorch Geometric官方教程,改編出從GCN基礎到點云分類的實操案例,每個例子都帶可運行代碼和錯誤分析,降低上手難度。
主要貢獻:讓非CUDA硬件"能用、好用"
-
降低硬件門檻:以前只有CUDA顯卡能高效運行幾何學習模型,現在研究人員用Gaudi HPU也能做實驗,節省成本且拓寬硬件選擇。
-
提供一站式資源:公開的GitHub倉庫整合了工具代碼、教程和故障解決方案,相當于一個"幾何學習跨平臺加油站",地址:https://github.com/NAVER-INTEL-Co-Lab/gaudi-geometric-learning。
-
推動跨平臺發展:為幾何學習算法在非CUDA硬件上的優化提供了基礎,就像為軟件跨平臺運行鋪好了道路。
總結
解決的主要問題
論文解決了Intel Gaudi-v2 HPU與PyTorch幾何學習框架的兼容性問題,通過重寫核心操作、提供實踐教程,讓非CUDA硬件能順利運行圖神經網絡模型,填補了現有庫對新興加速器支持的空白。
主要成果
- 開發了Gaudi HPU兼容的散射、稀疏等操作工具集;
- 整理了16個教程和11個案例,覆蓋從基礎到復雜的幾何學習任務;
- 建立公開倉庫,分享移植經驗和故障解決方案,助力非CUDA硬件上的幾何學習研究。
一段話總結
本文聚焦于在Intel Gaudi-v2 HPU等非CUDA處理單元上運行基于PyTorch的幾何學習框架的實踐,介紹了開發的核心實用工具(如散射、稀疏索引、k近鄰等操作),提供了十六個引導教程和十一個真實案例,包含遇到的故障診斷分析及詳細解決方法,相關資源整合在公開的GitHub倉庫中,為研究人員在非CUDA硬件上實驗幾何學習算法和模型降低了門檻,為進一步優化和跨平臺可移植性奠定了基礎。
思維導圖
詳細總結
一、研究背景與目標
- 幾何學習的重要性:幾何學習已成為建模非歐幾里得數據(尤其是圖結構數據)的強大范式,應用涵蓋社交網絡、分子結構、知識圖譜和推薦系統等領域。
- 硬件現狀:Nvidia支持CUDA的GPU在深度學習硬件領域占主導地位,但英特爾的Gaudi HPU等新興加速器在性能和能效方面具有競爭力。
- 面臨挑戰:使用此類非CUDA處理單元需要大量的工程工作和新穎的軟件適配,例如稀疏操作不支持、關鍵內核的反向傳播問題等。
- 研究目標:將基于PyTorch的幾何學習框架移植到Gaudi-v2 HPU,降低研究人員在非CUDA硬件上實驗幾何學習算法和模型的門檻。
二、核心貢獻
- 開發核心實用工具:恢復了Gaudi-v2 HPU上的基本操作,如散射、稀疏索引、k-最近鄰等。
- 提供豐富資源:整合了16個引導教程和11個真實世界的例子,并對遇到的故障進行了診斷分析,提供了詳細的解決方法。
- 建立公開倉庫:將所有經驗整合到一個公開的GitHub存儲庫中,為進一步優化和跨平臺可移植性提供了基礎。
三、GitHub存儲庫概述
部分 | 內容詳情 |
---|---|
文檔 | 包含用戶友好的README,包括環境設置說明、在Gaudi上啟動作業的指導以及相關外部材料的鏈接 |
核心實用工具 | 一套輔助模塊,旨在彌補現有庫的不足,支持在Gaudi上進行散射、稀疏索引和鄰域計算等操作,無需底層編程 |
教程和例子 | 包含分步教程和實際例子,涵蓋從基本圖卷積網絡(GCN)構建到鏈接預測和點云分類等任務,每個都有可運行的代碼和對在Gaudi-v2上遇到的挑戰的分析 |
四、核心實用工具
- 散射操作
- 挑戰:torch-scatter庫提供了CPU和CUDA后端的實現,但缺乏對Gaudi HPU的原生支持。
- 解決方案:使用純PyTorch函數重新實現散射操作,確保與Gaudi-v2 HPU兼容。
- 稀疏操作
- 挑戰:torch-sparse庫提供了CPU和CUDA后端的實現,但缺乏對Gaudi HPU的原生支持。
- 解決方案:將幾個稀疏張量操作轉換為等效的散射操作,例如稀疏矩陣和密集矩陣之間的乘法可以通過基于散射的聚合重新實現。
- 高級圖操作
- 挑戰:k-最近鄰(k-NN)搜索和METIS分區等高級圖操作缺乏對Gaudi HPU的支持。
- 解決方案:提供在CPU上執行這些圖實用程序的指南和實現,然后將結果傳輸到Gaudi HPU進行進一步處理。
五、教程和例子
- 教程
- GCN和GAT管道:涵蓋基本的圖卷積和注意力機制。
- 嵌入和聚合函數:說明如何在Gaudi上高效實現和執行節點嵌入和聚合函數。
- 圖池化和讀出:演示各種用于圖級表示學習和總結的技術。
- DeepWalk和Node2Vec實踐:通過隨機游走介紹表示學習,強調用于下游任務的節點嵌入。
- 例子
- 節點分類:使用基準數據集的分步工作流程,包括訓練和評估程序。
- 鏈接預測:演示如何使用GNN建模實體之間的關系,包括推薦等任務的實際工作流程。
- 用DGCNN進行點云分類:說明超越圖的幾何深度學習,特別是針對3D點云分類任務。
- 診斷分析
- 記錄了將現有PyTorch幾何教程和例子適配到Gaudi-v2 HPU時遇到的各種兼容性問題,包括不支持的操作、意外的運行時錯誤和細微的數值不匹配,并提供了系統的調試過程和解決方法。
六、結論
本文展示了將基于PyTorch的幾何學習工作流適配到Intel Gaudi-v2 HPU的努力,通過開發核心實用工具彌補了功能差距,提供的教程和例子展示了在非CUDA硬件上的端到端GNN訓練,所有資源整合在公開倉庫中,降低了研究人員在替代加速器上探索圖神經網絡的門檻。
關鍵問題
- 問題:在Gaudi-v2 HPU上運行PyTorch幾何學習框架面臨的主要挑戰是什么?
答案:主要挑戰包括現有庫(如torch-scatter、torch-sparse、torch-cluster)缺乏對Gaudi HPU的原生支持,導致散射操作、稀疏操作和高級圖操作(如k-NN搜索)無法直接使用,還存在反向傳播等問題。 - 問題:針對散射操作缺乏支持的問題,采取了什么解決方案?
答案:使用純PyTorch函數重新實現了散射操作,確保與Gaudi-v2 HPU兼容,無需依賴原生不支持的torch-scatter庫。 - 問題:GitHub存儲庫中包含哪些主要部分來支持在Gaudi-v2 HPU上的幾何學習?
答案:存儲庫包含三個主要部分,分別是文檔(提供環境設置、作業啟動等指導)、核心實用工具(實現散射、稀疏等操作)、教程和例子(涵蓋從基礎到實際應用的各種任務及代碼和分析)。