DGL-NPU：昇腾原生加速，高效赋能图神经网络开发

来源：市场资讯

（来源：华为计算）

AI for Science浪潮正持续驱动科研范式革新。依托图机器学习强大的拓扑结构解析与建模能力，蛋白质设计、分子属性预测、大分子相互作用仿真等生命科学、先进材料前沿场景实现研发效率大幅提升。

DGL（Deep Graph Library）作为行业主流图神经网络开源加速库，已深度落地各类AI4S研究场景，是全球科研团队开展生物大分子建模、分子筛选、药物仿真等实验的核心技术底座。长期以来，DGL原生框架暂未提供昇腾NPU适配能力，图神经网络模型难以直接在昇腾AI软硬件平台高效运行，限制了昇腾AI软硬件平台在图科学计算领域的应用落地。

为进一步拓宽图神经网络算力生态适配边界，充分释放昇腾AI基础软硬件在基础科研场景的算力效能，华为联合北京邮电大学石川教授团队，正式启动DGL适配昇腾专项攻坚项目。项目将完成DGL框架与昇腾AI基础软硬件平台深度适配、算子调优与稳定性打磨，实现各类图神经网络模型在昇腾NPU上高效、稳定、规模化部署运行，为AI4S前沿科研创新提供更强算力支撑。

DGL架构设计

DGL-NPU 的设计沿用 DGL 原有的分层架构，实现层面遵循开闭原则，只扩展，不修改，既保留了 DGL 原有的灵活性，也为昇腾 AI 基础软硬件平台提供了统一、高效、可扩展的运行支持，为后续合入上游代码仓奠定基础。

整体架构：用户Python层、DGLGraph对象层、FFI调用层、DGL C++ Runtime、Backend抽象层以及底层硬件层。

图数据管理：DGLGraph将图结构数据与特征数据解耦处理：

数据交换：DGL-NPU通过DLPack支持零拷贝数据交换，降低不同运行时与深度学习框架之间的数据搬运开销。

设备扩展：在原有CPUDeviceAPI的基础上扩展AscendDeviceAPI，使DGL C++ Runtime能够统一管理CPU与NPU。

核心算子：

图 1 DGL-NPU架构图

适配进展

核心图算子支持

围绕图神经网络中的高频计算模式，北京邮电大学石川教授团队实现了多类核心算子，为上层模型提供基础支撑。

目前DGL-NPU已支持：

• SpMM (sum, mean, max, min)

• BSpMM (sum, mean, max, min)

• SegmentReduce (sum, mean, max, min)

• RandomWalk

通过对sum、mean、max、min等常见规约方式的支持，DGL-NPU能够覆盖多种图模型中的核心计算需求。

典型模型适配

基于上述算子能力，北京邮电大学石川教授团队实现了多个典型图学习模型的支持，包括：

• LightGCN：常用于推荐系统场景，其核心计算依赖图上的邻居传播与聚合；

• GraphSAGE：面向大规模图的归纳式表示学习，对采样和聚合能力有较高要求；

• HEREC：则面向异构信息网络，通过随机游走等机制学习节点表示。

这些模型的支持，验证了DGL-NPU在推荐、节点分类、异构图表示学习等典型任务中的可用性，也为更多图神经网络模型迁移到昇腾AI基础软硬件平台奠定了基础。

分布式训推支持

基于上述算子能力，华为团队基于HCCL实现将大图切分成多个子图进行分布式训推的能力，包括图分区创建与索引、基于共享内存的分布式数据访问，以及针对大规模嵌入的参数更新等能力，实现了COOToCSR、CSRToCOO、CSRSliceRows、CSRSliceMatrix、CSRGetRowNnz、CSRGetData、BinaryElewise和IndexSelect算子的适配，基于Cora数据集做了验证。

实战案例

OpenPOM模型基于DGL-NPU单卡运行示例

OpenPOM 的建立基于一篇2023年发表于《科学》杂志的里程碑式论文《A principal odor map unifies diverse tasks in human olfactory perception》，是首个能像人类一样可靠地描述未知分子气味的计算模型，平均 ROC-AUC 达到了 0.8872，排名第一，为气味数字化奠定了基础，后续研究也在此基础上展开。

# 安装deepchempip install --pre deepchem# 下载dgl源码并编译git clone --recursive https://github.com/BUPT-GAMMA/dgl-ascend.gitcd dgl-ascend/scriptbash build_dgl_ascend.shcd ../python# 安装dglpip install -e .# 回到主目录cd ../..# 下载openpom源码git clone https://github.com/lizengyong/openpom.gitcd openpom# 用cpu执行预测python predict_odors_cli.py "CCCCCCCCCC" -n 10 --device cpu# 用npu执行预测python predict_odors_cli.py "CCCCCCCCCC" -n 10 --device npu:0

CPU预测结果

NPU预测结果

正癸烷（n-Decane）分子气味预测结果对比

基于NPU预测正癸烷（n-Decane）的味道，结果与基于CPU一致，速度是CPU的20倍。

GCN模型基于DGL-NPU 单机多卡分布式训练示例

图卷积网络（GCN）是一种专为图结构数据设计的深度学习模型，核心思想是通过聚合节点自身及其邻居的特征信息，来生成蕴含局部拓扑结构的节点表示。它突破了传统卷积网络只能处理图像、文本等规则网格数据的局限，使得对社交网络、分子结构或知识图谱等非欧几里得数据进行高效的端到端学习成为可能，为关系推理和结构化数据建模提供了强大的通用范式，不仅大幅推动了药物研发、物理模拟等科学计算领域的发展，更成为连接复杂系统数据与现代人工智能算法的核心桥梁。

# 下载dgl源码并编译git clone --recursive https://github.com/lizengyong/dgl-ascend.gitcd dgl-ascend/scriptbash build_dgl_ascend.shcd ../pythonpip install -e .cd ../tests/ascend# 单卡训练torchrun --nproc_per_node=1 --master_port=29502 test_dist_gcn_converge_npu.py --dataset ogbn-arxiv --mode update_all --batch-size 64 --epochs 30 --lr 0.002# 单机2卡分布式训练torchrun --nproc_per_node=2 --master_port=29502 test_dist_gcn_converge_npu.py --dataset ogbn-arxiv --mode update_all --batch-size 64 --epochs 30 --lr 0.002# 单机8卡分布式训练torchrun --nproc_per_node=8 --master_port=29502 test_dist_gcn_converge_npu.py --dataset ogbn-arxiv --mode update_all --batch-size 64 --epochs 30 --lr 0.002

单卡30 epoch训练结果

单机2卡30 epoch训练结果

单机8卡30 epoch训练结果

基于2层GCN模型和ogbn-arxiv数据集30 epoch训练结果对比

基于2层GCN模型和ogbn-arxiv数据集（169343节点，2484941条边，90941训练节点）进行30 epoch 训练，单卡、单机2卡和单机8卡均loss持续收敛、TrAcc趋势上升和ValAcc趋势上升，单机多卡分布式训练效率比单卡提升明显。

总结与展望

华为将持续联合北京邮电大学石川教授团队，长期支持DGL的适配和优化工作，推动DGL-NPU在更多实际业务场景中的落地应用，也欢迎对于图机器学习有兴趣的同学，一起参与到项目中或者积极提出模型适配、使用反馈。

立即开始--共建AI4S昇腾生态

开源代码仓地址：

GitHub地址：

https://github.com/BUPT-GAMMA/dgl-ascend

GitCode地址：https://gitcode.com/AI4Science/dgl-ascend