pytorch-adapter：让 PyTorch 模型“无缝”跑在昇腾 NPU 上

pytorch-adapter：让 PyTorch 模型“无缝”跑在昇腾 NPU 上

之前帮朋友看 PyTorch 模型适配 CANN 的代码，发现他手写了很多适配层——把自己的 MyModel 一层层翻译成 AscendCL 接口，光写适配层就写了 2,000 行。

我告诉他：不用手写，用 pytorch-adapter 就行。 这个适配器能自动把 PyTorch 模型跑在 NPU 上，不用改模型代码（或者只改几行）。

环境准备：装 pytorch-adapter 和依赖

在拆 pytorch-adapter 的用法之前，先把环境装好。不然后面跑代码报“模块找不到”，又得回头查。

第1步：装 CANN（必备）

pytorch-adapter 依赖 CANN 的 AscendCL 接口，得先装 CANN。推荐装 CANN 8.0+（对 PyTorch 2.x 支持更好）。

# 检查 CANN 是否装好npu-smi info

如果看到 NPU 设备信息，说明 CANN 装好了。

⚠️踩坑预警：CANN 版本跟 pytorch-adapter 版本要对应。CANN 8.0 得配 pytorch-adapter v3.x，配错了模型跑不起来。

第2步：装 PyTorch（推荐 2.1+）

pytorch-adapter 支持 PyTorch 1.13-2.1，推荐用 PyTorch 2.1（性能更好）。

# 装 PyTorch 2.1（CPU 版本就行，adapter 会替换成 NPU 版本）pipinstalltorch==2.1.0torchvision==0.16.0torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cpu

⚠️踩坑预警：装的是 CPU 版本的 PyTorch（不是 CUDA 版本）。pytorch-adapter 会替换 PyTorch 的后端（从 CUDA 换成 NPU），如果你装了 CUDA 版本，会冲突。

第3步：装 pytorch-adapter

# 方法1：pip 安装（推荐）pipinstallpytorch-npu==3.0.0# 对应 CANN 8.0# 方法2：源码编译（如果你想改 adapter 的代码）gitclone https://atomgit.com/cann/pytorch-adapter.gitcdpytorch-adapter&&gitcheckout v3.0# 对应 CANN 8.0python setup.pyinstall

装完后，Python 里能import torch并且torch.npu.is_available()返回True，就说明装好了。

逐步实现：让 PyTorch 模型跑在 NPU 上

第1步：把模型搬到 NPU 上（.npu()）

pytorch-adapter 提供了.npu()方法（类似 PyTorch 的.cuda()），把模型/张量搬到 NPU 上。

importtorchimporttorch.nnasnn# 1. 定义模型（跟普通 PyTorch 模型一样）classMyModel(nn.Module):def__init__(self):super().__init__()self.fc1=nn.Linear(784,512)self.fc2=nn.Linear(512,10)defforward(self,x):x=torch.relu(self.fc1(x))x=self.fc2(x)returnx model=MyModel()# 2. 把模型搬到 NPU 上（关键！）model=model.npu()# 类似 .cuda()# 3. 把输入数据搬到 NPU 上input_data=torch.randn(32,784).npu()# 类似 .cuda()# 4. 前向计算（自动在 NPU 上跑）output=model(input_data)print(output.device)# 输出：npu:0

关键点：

• .npu()是 pytorch-adapter 自动添加的方法（不用改 PyTorch 源码）
• 模型搬到 NPU 后，所有计算都自动在 NPU 上跑（不用手动调算子）
• ⚠️ 别忘了把输入数据也搬到 NPU 上（input_data = input_data.npu()）。如果模型在 NPU 上，输入数据在 CPU 上，会报“设备不匹配”错误。

第2步：训练模型（自动在 NPU 上算梯度）

pytorch-adapter 自动替换了 PyTorch 的后端（从 CUDA 换成 NPU），所以训练代码不用改（跟 GPU 训练一模一样）。

importtorchimporttorch.nnasnnimporttorch.optimasoptim# 1. 定义模型 + 搬到 NPUmodel=MyModel().npu()# 2. 定义损失函数 + 优化器（跟 GPU 训练一模一样）criterion=nn.CrossEntropyLoss()optimizer=optim.Adam(model.parameters(),lr=0.001)# 3. 训练循环（跟 GPU 训练一模一样）forepochinrange(10):forbatch_idx,(data,target)inenumerate(train_loader):# 把数据搬到 NPUdata,target=data.npu(),target.npu()# 前向计算output=model(data)loss=criterion(output,target)# 反向传播（自动在 NPU 上算梯度）optimizer.zero_grad()loss.backward()optimizer.step()ifbatch_idx%100==0:print(f'Epoch{epoch}, Batch{batch_idx}, Loss{loss.item()}')

关键点：

• loss.backward()会自动在 NPU 上算梯度（pytorch-adapter 替换了 autograd 后端）
• 优化器会自动更新 NPU 上的参数（不用手动拷贝）
• ⚠️ 如果你用了自定义的 autograd 函数（torch.autograd.Function），得手动改（adapter 没法自动替换）。得把ctx.cuda()改成ctx.npu()。

第3步：保存/加载模型（NPU 格式）

训练完的模型，可以用 PyTorch 标准的保存/加载接口（adapter 自动处理了 NPU 格式）。

# 1. 保存模型（跟 GPU 训练一模一样）torch.save(model.state_dict(),'model_npu.pth')# 2. 加载模型（得先搬到 NPU 上）model=MyModel().npu()# 先搬到 NPUmodel.load_state_dict(torch.load('model_npu.pth'))model.eval()# 推理模式# 3. 推理input_data=torch.randn(1,784).npu()output=model(input_data)pred=output.argmax(dim=1)print(f'Prediction:{pred.item()}')

⚠️踩坑预警：保存的模型是 NPU 格式的（不是 CPU 格式）。如果你要拿到 CPU 上推理，得先搬到 CPU（model = model.cpu()），再保存。

性能调优：让模型跑得更快

调优点1：开启算子融合（Graph Mode）

pytorch-adapter 支持 Graph Mode（图模式），会自动融合算子（类似 GE 图引擎的功能），性能提升 20%-30%。

importtorch# 开启 Graph Mode（在模型推理前调用）torch.npu.enable_graph_mode(True)# 推理（会自动做算子融合）input_data=torch.randn(1,784).npu()output=model(input_data)

⚠️踩坑预警：Graph Mode 不是所有模型都支持。如果你的模型有动态控制流（if/for/while），Graph Mode 会报错。这时候得关掉 Graph Mode（torch.npu.enable_graph_mode(False)）。

调优点2：用 NPU 的加速库（ATB）

如果你用的是 Transformer 类模型（BERT/GPT/LLaMA），可以用 ATB（Ascend Transformer Boost）加速，性能提升 50%-100%。

fromatbimportTransformerInfer# 用 ATB 包装模型model=MyGPTModel().npu()infer=TransformerInfer(model,max_seq_len=2048,soc_type="Ascend910")# 推理（自动用 ATB 加速）input_ids=torch.tensor([[1,2,3,4,5]]).npu()output=infer.forward(input_ids)

调优点3：用混合精度训练（FP16）

NPU 的 FP16 算力是 FP32 的 2-3 倍。用混合精度训练，性能提升 50%-80%。

importtorch# 开启自动混合精度（跟 GPU 的 amp 用法一样）withtorch.npu.amp.autocast():output=model(input_data)loss=criterion(output,target)loss.backward()

踩坑实录

踩坑1：模型搬不到 NPU 上（.npu() 报错）

原因：PyTorch 版本跟 pytorch-adapter 不匹配。

解决方案：检查版本对应关系：

踩坑2：训练时梯度算不出来（loss.backward() 报错）

原因：用了自定义的 autograd 函数（torch.autograd.Function），adapter 没法自动替换。

解决方案：手动改自定义函数的代码，把ctx.cuda()改成ctx.npu()。

踩坑3：推理性能不达标（比 GPU 慢）

原因：没开 Graph Mode 或者没用 ATB 加速库。

解决方案：

• 开启 Graph Mode（torch.npu.enable_graph_mode(True)）
• 如果是 Transformer 模型，用 ATB 加速（TransformerInfer）

实践指引

1. 读 pytorch-adapter 源码：从torch/npu/__init__.py看起，理解.npu()是怎么注入的
2. 跑 pytorch-adapter 的示例：pytorch-adapter 仓库里有现成的示例（examples/目录）
3. 调 Graph Mode：如果你的模型推理性能不达标，试试开 Graph Mode
4. 用 ATB 加速 Transformer 模型：如果是 Transformer 类模型，用 ATB 加速，性能提升 50%-100%

仓库链接：
https://atomgit.com/cann/pytorch-adapter
https://atomgit.com/cann/ascend-transformer-boost
https://atomgit.com/cann/runtime