作为 Ascend for PyTorch 社区的核心组件，TorchNPU 是昇腾专为 PyTorch 打造的深度学习适配插件，使 PyTorch 框架能够直接调用昇腾 NPU，为开发者提供昇腾 AI 处理器的超强算力。

This repository is empty

TorchNPU

基于昇腾NPU的PyTorch适配插件

简体中文 | English

支持版本 · 安装 · 快速开始 · 版本配套 · 支持说明 · 贡献指南 · 文档 · 社区

version rust platform

简介

作为 Ascend for PyTorch 社区的核心组件，TorchNPU 是昇腾专为 PyTorch 打造的深度学习适配插件，使 PyTorch 框架能够直接调用昇腾 NPU，为开发者提供昇腾 AI 处理器的超强算力。

昇腾为基于华为昇腾 AI 处理器和软件的行业应用及服务提供全栈 AI 计算基础设施。您可以通过访问昇腾社区，了解关于昇腾的更多信息。

核心特性

TorchNPU 充分继承和复用上游 PyTorch 的大量成熟功能，并在此基础上针对昇腾 NPU 进行深度适配与优化，下图展示了 TorchNPU 的主要功能模块：

TorchNPU 主要模块介绍：

基础计算： 广泛支持 PyTorch 原生 API 及自定义 API，覆盖主流AI场景，提供一致性体验，支撑用户快速实现模型和算法。
分布式： 支持通过 FSDP2 加速分布式训练，核心计算 API 支持 Dtensor，支持 AllGather、AllReduce、AllToAll 等集合通信原语，和 Send、Recv 等点对点通信原语。
图模式： 通过“动态图捕获+静态图优化+高效代码生成”的方式显著加速模型训练和推理任务，并支持通过 NPUGraph 下沉执行，在 2.6.0 以上版本已支持。
调试调优： 支持 Profiling 分析计算、通信和内存使用，支持通过 WatchDog 实时监控分析通信异常。
TorchNPU Core： 支持虚拟内存管理降低内存碎片，在分布式场景支持跨流内存复用优化，通过 PrivateUse1 将算子和设备资源接入 PyTorch。

支持版本

当前分支以单一源码适配多个 PyTorch 版本，所支持的 PyTorch 版本清单声明在 version.txt 中，该文件是版本支持的唯一来源，随版本适配同步更新。

已发布软件包的版本配套关系（TorchNPU / PyTorch / CANN / Python）请参见版本配套。

安装

软件包安装

以安装 TorchNPU 2.10.0.post2 版本为例，请按照以下命令进行安装。下载其他版本的方式请参见社区下载页面 TorchNPU 下载。

安装 CANN

安装 9.0.0 版本的 CANN，具体步骤请参见 CANN 安装指南。

安装 PyTorch

执行以下命令安装 PyTorch 2.10.0：

pip install torch==2.10.0 --index-url https://download.pytorch.org/whl/cpu

安装 TorchNPU

执行以下命令安装 TorchNPU 2.10.0.post2：

pip install torch-npu==2.10.0.post2

源码编译安装

编译TorchNPU的详细步骤请参见源码安装。

快速开始

初始化环境

# 默认路径，请根据实际安装位置修改
source /usr/local/Ascend/ascend-toolkit/set_env.sh

运行示例

import torch
import torch_npu

x = torch.randn(2, 2).npu()
y = torch.randn(2, 2).npu()
z = x.mm(y)
print(z)

从 TorchNPU 2.5.1 版本起，import torch_npu 不再是强制要求（会自动注册），但建议在代码中显式导入以确保设备初始化。

输出如下类似信息说明运行成功:

tensor([[-0.0515,  0.3664],
        [-0.1258, -0.5425]], device='npu:0')

完整的模型迁移和训练教程，请参见快速入门。

社区交流

Ascend for PyTorch 社区由多个 Special Interest Groups（SIGs）组成，每个 SIG 负责特定技术领域的开发、维护和社区协作。以下是当前所有 SIG 的列表，点击对应链接可查看详细说明。

SIG名称	简要描述	链接
Core SIG	聚焦于昇腾 NPU 平台上的 PyTorch 核心适配层开发，负责 `TorchNPU` 扩展库及其算子插件 `OpPlugin` 的设计、实现与维护。	🔗 了解更多
Distributed SIG	致力于在昇腾 NPU 硬件底座上，围绕 PyTorch 分布式训练框架（torch.distributed）构建高效、易用、可扩展的并行训练能力，为大语言模型、多模态模型和强化学习等场景提供极致性能体验。	🔗 了解更多
Graph Mode SIG	聚焦于 Dynamo、Inductor、NPUGraph等核心技术，旨在通过自动化的图捕捉与编译优化技术，弥合“易用性”与“高性能”之间的鸿沟。	🔗 了解更多
Usability SIG	致力于推动Ascend for PyTorch易用性体验提升，包含文档、教程、案例等。	🔗 了解更多

每个SIG都有自己的例会、邮件列表和贡献指南。点击对应的SIG链接可查看详细联系方式、工作目标及参与方式。欢迎大家为社区做贡献。如果有任何疑问或建议，请提交 GitCode Issues，我们会尽快回复。感谢您的支持。

安全声明

TorchNPU 的系统安全加固、运行用户建议和文件权限控制等内容，请参见安全声明。

免责声明

本插件仅供调试和开发使用。使用者需自行保证输入命令行的安全性，并对使用过程中产生的数据做好权限控制。使用本插件即表示您同意并接受以上声明。

License

TorchNPU 的使用许可证，请参见 LICENSE 文件。 TorchNPU 资料文档的使用许可证，请参见 LICENSE 文件。

致谢

感谢来自社区的每一个 PR，欢迎开发者向 TorchNPU 插件贡献代码！

pytorch:基于昇腾NPU的PyTorch框架适配插件项目