PyTorch + NPU（torch_npu）¶
约 784 个字 26 行代码 预计阅读时间 3 分钟
torch_npu 是昇腾官方维护的 PyTorch 适配插件（Ascend Extension for PyTorch）。装上它之后，你熟悉的 PyTorch 训练/推理代码几乎不用改，把设备从 cuda 换成 npu 就能跑在昇腾上。
安装¶
前置条件：已经按环境搭建装好 NPU 驱动、CANN Toolkit 与 Kernels 包，并 source /usr/local/Ascend/ascend-toolkit/set_env.sh。
第一步：检查版本配套。 torch_npu、PyTorch、CANN 三者的版本必须严格配套，例如（节选自官方版本配套表，写作时点为 2026 年 8 月，最新配套以官方为准）：
| CANN | PyTorch | torch_npu | 备注 |
|---|---|---|---|
| 9.1.0 | 2.7.1 | 2.7.1.post8 | 本指南示例采用 |
| 9.1.0 | 2.12.0 | 2.12.0 | 写作时点的最新组合 |
| 8.2.RC1 | 2.6.0 | 2.6.0 | 历史版本举例 |
第二步：安装。 官方提供 whl 包，按 CPU 架构（aarch64/x86_64）与 Python 版本区分，安装前先在配套表中确认自己的环境被支持，然后 pip 安装即可。下文以 CANN 9.1.0 + PyTorch 2.7.1 + torch_npu 2.7.1.post8（支持 Python 3.9–3.13）为例，该组合同时是下一章 triton-ascend 的官方推荐环境：
# aarch64（鲲鹏/ARM 服务器）
pip install torch==2.7.1
# x86_64 服务器则安装 CPU 版（+cpu 版本仅在 PyTorch 官方源，需指定 index-url）
pip install torch==2.7.1+cpu --index-url https://download.pytorch.org/whl/cpu
pip install pyyaml setuptools
pip install torch-npu==2.7.1.post8
源码仓库在 Gitee: ascend/pytorch（已同步迁移至 GitCode），里面始终维护着最新的版本配套表。
第三步：验证。
source /usr/local/Ascend/ascend-toolkit/set_env.sh
python -c "import torch, torch_npu; print(torch.randn(2, 2).npu())"
基本使用¶
torch_npu 尽量保持了与 CUDA 版 PyTorch 相同的 API 风格：
1 2 3 4 5 6 7 8 9 10 11 12 | |
迁移现有代码¶
手动迁移通常只需全局替换：
cuda→npu（torch.device("cuda")、.cuda()、device="cuda"等）torch.cuda.xxx→torch_npu.npu.xxx（如torch.cuda.synchronize()）
自动迁移则更省事，在训练脚本最开头加一行：
from torch_npu.contrib import transfer_to_npu
它会拦截后续代码中对 torch.cuda 的调用并转发到 NPU，device="cuda"、torch.cuda.amp 等都能正确映射，适合快速跑通开源项目。基于 HuggingFace accelerate 的项目在装好 torch_npu 后，Accelerator().prepare(...) 会自动识别并使用 NPU 后端（注意 transfer_to_npu 与 accelerate 同时使用可能干扰后端检测，建议二选一）。
迁移报错怎么排查
先确认三件事：① set_env.sh 是否 source；② CANN 与 torch_npu 版本是否配套；③ 代码里是否绕过了迁移层直接调用了 CUDA 专有 API（如 pycuda、torch.cuda 底层接口）。
分布式与混合精度¶
多卡训练把通信后端从 NCCL 换成昇腾的 HCCL 即可，其余（DDP、torchrun）保持不变：
torchrun --nproc_per_node=8 train.py
torch.distributed.init_process_group(backend="hccl")
混合精度方面，Ascend 910B 支持 fp16 与 bf16，用 PyTorch 标准接口即可：
with torch.autocast("npu", dtype=torch.bfloat16):
out = model(input)
性能分析¶
torch_npu.profiler 与 torch.profiler 接口一致，可以导出 chrome trace 看算子耗时，定位搬运与计算的重叠情况；设备层面用 npu-smi info watch 观察利用率。
学习&拓展：¶
- 跑通本文“验证”小节的两行代码。
- 找一段你写过的（或开源的）CUDA 单卡训练脚本，用
transfer_to_npu迁移到 NPU 上并完整训练几个 epoch。 - 对比同一模型在 GPU 与 NPU 上的收敛速度与吞吐，尝试解释差异来源（算子覆盖度、混合精度行为等）。
- 用双卡跑一次 DDP，确认
backend="hccl"下梯度能正确同步。 - 用
torch_npu.profiler导出 trace，找出耗时最长的算子。