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 不匹配。
解决方案:检查版本对应关系:
| PyTorch 版本 | pytorch-adapter 版本 |
|---|---|
| 1.13 | v2.0.x |
| 2.0 | v3.0.x |
| 2.1 | v3.1.x |
踩坑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)
实践指引
- 读 pytorch-adapter 源码:从
torch/npu/__init__.py看起,理解.npu()是怎么注入的 - 跑 pytorch-adapter 的示例:pytorch-adapter 仓库里有现成的示例(
examples/目录) - 调 Graph Mode:如果你的模型推理性能不达标,试试开 Graph Mode
- 用 ATB 加速 Transformer 模型:如果是 Transformer 类模型,用 ATB 加速,性能提升 50%-100%
仓库链接:
https://atomgit.com/cann/pytorch-adapter
https://atomgit.com/cann/ascend-transformer-boost
https://atomgit.com/cann/runtime