PyTorch 模型迁移昇腾 NPU 实战：torch_npu 适配与性能调优

昇腾生态里，应用开发者最常见的任务不是重写模型，而是把已有的 PyTorch 代码迁移到 NPU 上跑起来，并且跑得够快。好消息是，torch_npu 作为 PyTorch 的昇腾适配插件，让这件事的代码改动量通常只有几十行。

坏消息是，"能跑"和"跑得好"之间，还隔着精度对齐、混合精度、多卡通信、算子回退、性能瓶颈排查一整套工程问题。本文按"先正确、后性能"的顺序，把这条路径完整走一遍。

一、迁移全景：三个原则

开始之前，先把原则立好，可以避免大量返工：

1. 先正确、后性能：一开始不要开各种加速开关，先保证 loss 曲线对得上；
2. 先单卡、后多卡：单卡精度没对齐就上分布式，问题会呈指数级放大；
3. 先小规模、后全量：用固定随机种子 + 小数据集对拍，比对全量训练更高效。

版本配套是迁移的第一道门槛。PyTorch、torch_npu、CANN、驱动固件四者必须严格配套，版本错位导致的报错往往非常隐晦。安装前务必查阅官方版本配套表。

二、环境检查：三行命令

# 1. 确认设备侧状态
npu-smi info

# 2. 确认 Python 侧适配
python -c "
import torch
import torch_npu
print('torch:', torch.__version__)
print('npu available:', torch.npu.is_available())
print('npu count:', torch.npu.device_count())
print('device name:', torch.npu.get_device_name(0))
"

如果第二步失败，按"驱动 → CANN → torch_npu → PyTorch"的顺序逐层排查，不要直接怀疑模型代码。

三、最小代码改动：从 cuda 到 npu

迁移的核心动作只有两个：引入 torch_npu、把设备字符串换掉。

# 迁移前（GPU）
import torch

device = torch.device("cuda:0")
model = model.cuda()
x = x.cuda()
scaler = torch.cuda.amp.GradScaler()

torch.cuda.synchronize()

# 迁移后（NPU）
import torch
import torch_npu  # 必须导入，注册 npu 后端

device = torch.device("npu:0")
model = model.to(device)
x = x.to(device)
scaler = torch.npu.amp.GradScaler()

torch.npu.synchronize()

常用 API 的对应关系：

如果代码里到处是硬编码的 .cuda()，建议先做一次简单重构，统一走一个 device 变量，后续换卡成本会低很多。

3.1 DataLoader 注意事项

• pin_memory=True 依旧有效，配合 non_blocking=True 可以减少 H2D 拷贝等待；
• num_workers 不是越大越好，昇腾服务器通常 CPU 核多，但过多的 worker 会造成 IPC 开销，建议从 4～8 起步做对比实验；
• 自定义 Dataset 中若有 CUDA 相关调用（如 torch.cuda.current_device()），需要同步替换。

四、混合精度与融合优化器

昇腾对 FP16/BF16 的支持比较完善。开启混合精度的方式与 CUDA 侧几乎一致：

import torch
import torch_npu

model = model.to("npu:0")
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = torch.npu.amp.GradScaler()

for images, labels in dataloader:
    images = images.to("npu:0", non_blocking=True)
    labels = labels.to("npu:0", non_blocking=True)

    optimizer.zero_grad(set_to_none=True)
    with torch.npu.amp.autocast():
        loss = criterion(model(images), labels)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

此外，torch_npu 提供了一批融合优化器，把参数更新中的多个小算子合并，既省显存又减少下发开销：

from torch_npu.optim import NpuFusedAdamW

optimizer = NpuFusedAdamW(
    model.parameters(),
    lr=1e-4,
    betas=(0.9, 0.999),
    eps=1e-8,
    weight_decay=0.01,
)

对训练大模型的团队来说，这类融合优化器通常是"免费的加速"，建议在精度对齐之后优先尝试。

五、多卡分布式训练

昇腾多卡的通信后端是 HCCL。代码结构基本不变：

import os

import torch
import torch.distributed as dist
import torch_npu

def main():
    dist.init_process_group(backend="hccl")

    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"npu:{local_rank}")
    torch.npu.set_device(device)

    model = build_model().to(device)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

    # ... 训练循环 ...

启动方式与 torchrun 一致：

torchrun --nproc_per_node=8 --master_port=29500 train.py

多机训练时还需要注意：

• 各节点 MASTER_ADDR、MASTER_PORT 一致，防火墙放通对应端口；
• 多网卡机器要通过环境变量（如 HCCL_SOCKET_IFNAME）指定通信网卡，选错网卡是最常见的"卡住不动"原因；
• 先用 2 机 2 卡做连通性验证，再扩到全规模。

六、精度对齐：怎么确认"算得对"

迁移后第一件事不是看速度，而是看 loss。建议按下面的顺序排查：

1. 固定随机种子、固定数据顺序，在 CPU/GPU 上产出一份 baseline（前 100～500 step 的 loss）；
2. NPU 上跑同样的配置，对比 loss 曲线。允许微小数值差异，但趋势必须一致；
3. 如果差异明显，先关闭 AMP 用 FP32 复测。很多"精度问题"其实是混合精度溢出；
4. 检查是否有算子回落到 CPU。回退不仅影响性能，数值实现也可能不同；
5. 用单算子对拍：把可疑层单独抽出来，输入同一份数据，比对输出；
6. 检查 loss scale 是否频繁触发缩放，必要时改用动态 loss scale 或 BF16。