2.4 torch_npu：把 PyTorch 模型迁移到昇腾 NPU

现实中的 AI Infra 工程师，九成九的日常工作不是「从零写框架」，而是「让现有代码跑起来」。PyTorch 已经是训练的事实标准，昇腾要接入生态，最务实的路径就是 torch_npu——一个让 PyTorch 代码几乎不改就能跑在昇腾 NPU 上的适配插件。这一章我们讲清楚它怎么工作、怎么用。

📑 目录

1. torch_npu 是什么
2. 原理：NPUDevice 与算子派发
3. 安装与环境准备
4. 最小迁移：一个完整的训练示例
5. 常见迁移清单与注意事项
6. 遇到不支持的算子怎么办
7. torch_npu 与 CUDA 代码的对照
总结
延伸阅读

1. torch_npu 是什么

torch_npu（TorchNPU）是 PyTorch 在昇腾 NPU 上的适配插件，由华为昇思（MindSpore）团队与社区共同维护。

一句话版本：

装上 torch_npu 后，PyTorch 多了一个叫 npu 的设备类型；把模型和数据 .to('npu') 过去，原来的训练/推理代码就能跑在昇腾芯片上。

它扮演的角色 ≈ PyTorch 官方 torch.cuda 之于 NVIDIA GPU：torch.cuda 是「PyTorch ↔ GPU」的桥，torch_npu 是「PyTorch ↔ 昇腾 NPU」的桥。

官方支持矩阵（以 vLLM-Ascend 官方 README 为准，2026 年当前推荐组合）为：

CANN 9.1.0
PyTorch 2.10.0
TorchNPU 2.10.0.post4

版本组合随 CANN / vLLM-Ascend 发布节奏变动，务必以昇腾社区「安装指南」的配套表与 vLLM-Ascend 官方 README 为准。

2. 原理：NPUDevice 与算子派发

torch_npu 的核心是「用 PyTorch 的算子派发机制，把算子路由到昇腾实现」：

引入 torch_npu 后，PyTorch 注册了一个新的设备类型 npu。
当你执行 tensor.to('npu') 或 model.to('npu')，张量和参数就被放到 NPU 设备上。
前向/反向中的每个算子（Conv、MatMul、LayerNorm……）触发 PyTorch 的设备分派（device dispatch）：算子被路由到 torch_npu 提供的 NPU 实现。
torch_npu 的实现内部调用 CANN 的 aclnn 算子库或昇腾高性能算子，最终在 AI Core 上执行。

PyTorch 代码 ──▶ torch_npu（NPU device 实现）──▶ aclnn/昇腾算子库 ──▶ CANN Runtime ──▶ 昇腾芯片

这意味着：只要某个算子在 torch_npu 里有对应实现，你的代码一行都不用改；没实现的算子，才会报错或回退。torch_npu 长期在补齐算子和优化性能两条线上迭代，这也是昇腾生态成熟度的晴雨表。

3. 安装与环境准备

典型安装步骤（以 CANN 已装好为前提，具体以官方文档为准）：

# 1. 确认 CANN 环境（source 环境脚本）
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 2. 确认 PyTorch 版本与 CANN 匹配
pip list | grep torch

# 3. 安装 torch_npu（与 PyTorch 版本对应的 wheel）
pip install torch-npu

# 4. 验证安装：能打印出 NPU 信息即成功
python -c "import torch, torch_npu; print(torch.npu.is_available(), torch.npu.device_count())"

昇腾官方也提供「昇腾 AI 容器镜像」，内含驱动/固件/CANN/框架全栈，省去逐项安装的麻烦——推荐生产环境使用。

4. 最小迁移：一个完整的训练示例

以经典的 PyTorch 训练循环为例，对比「GPU 版」和「昇腾 NPU 版」：

import torch
import torch.nn as nn
import torch_npu  # 关键：引入后 torch 才能识别 npu 设备

device = 'npu:0' if torch.npu.is_available() else 'cpu'

class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(16, 8)

    def forward(self, x):
        return self.fc(x)

model = TinyNet().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()

for step in range(10):
    x = torch.randn(32, 16, device=device)   # 直接在 npu 上建张量
    y = torch.randn(32, 8, device=device)
    pred = model(x)
    loss = loss_fn(pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if step % 5 == 0:
        print(f"step={step} loss={loss.item():.4f}")

迁移改动清单（对比原本 CUDA 代码）：

CUDA 写法	NPU 写法
`import torch`	`import torch, torch_npu`
`device = 'cuda'`	`device = 'npu:0'`
`torch.cuda.is_available()`	`torch.npu.is_available()`
`.to('cuda')`	`.to('npu')`
`torch.cuda.xxx` API	`torch.npu.xxx` API

其余（模型定义、Dataset、DataLoader、优化器、训练循环）完全不变。这就是 torch_npu 最大的价值：迁移成本极低。

5. 常见迁移清单与注意事项

设备字符串：cuda → npu，多卡时 npu:0/npu:1。
torch.cuda.* 相关 API：检查代码里是否有 torch.cuda.synchronize()、torch.cuda.memory_summary() 等，需替换为 torch.npu.*。
DataLoader：pin_memory=True 在 NPU 上通常不需要（对应 CUDA pinned memory 优化），可保留也可关闭。
随机种子：torch.manual_seed 通用；如需设备级种子用 torch.npu.manual_seed_all。
分布式：torch.distributed 支持通过 HCCL 后端初始化（详见第 4 章），替换 nccl 为 hccl。
混合精度：torch.cuda.amp 对应昇腾场景常用 torch.npu.amp（autocast/GradScaler 都有 NPU 版本）。

6. 遇到不支持的算子怎么办

这是实际迁移中最常见的问题。处理顺序推荐：

升级 torch_npu / CANN：很多算子缺口随版本迭代被补齐，先排除「版本太旧」。
查算子支持矩阵：昇腾社区/官方文档有算子的支持列表，确认该算子是否已适配。
降级到 CPU：将个别不受支持的算子移到 CPU 执行（tensor.cpu() 计算后再搬回 NPU），作为临时方案。
算子融合/替换：用等价的组合算子替代（如把 F.interpolate 换成可支持的实现）。
自己写算子：用 Ascend C 自定义算子，注册后供 torch_npu 调用——这是第 3 章的内容，也是 AI Infra 工程师在昇腾上最核心的增值技能。

7. torch_npu 与 CUDA 代码的对照

能力	CUDA（GPU）	torch_npu（NPU）
设备类型	`cuda`	`npu`
算子派发	torch 派发到 CUDA kernels	torch 派发到 aclnn/昇腾算子库
后端通信库	NCCL	HCCL
AMP 混合精度	`torch.cuda.amp`	`torch.npu.amp`
性能分析	Nsight / PyTorch Profiler	msprof / PyTorch Profiler（昇腾插件）
显存管理	`torch.cuda.memory_*`	`torch.npu.memory_*`

结论：把 torch_npu 理解成「昇腾版的 torch.cuda」，迁移就是「把 cuda 换成 npu」+「处理少量专属 API」的过程。

总结

torch_npu 让 PyTorch 模型几乎零成本迁移到昇腾 NPU，核心机制是「注册 npu 设备 + 算子派发到 CANN」。
迁移三步走：装对版本 → import torch_npu → cuda 换成 npu。
遇到不支持的算子，按「升级 → 查矩阵 → 降级 CPU → 算子替换 → 自研算子」的次序解决。
自研算子（Ascend C）是 AI Infra 工程师在昇腾生态里最能体现价值的技能，详见第 3 章。

搜索