在非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上的几何学习？
答案：存储库包含三个主要部分，分别是文档（提供环境设置、作业启动等指导）、核心实用工具（实现散射、稀疏等操作）、教程和例子（涵盖从基础到实际应用的各种任务及代码和分析）。