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.1pytorch-adapter 支持 PyTorch 1.13-2.1推荐用 PyTorch 2.1性能更好。# 装 PyTorch 2.1CPU 版本就行adapter 会替换成 NPU 版本pipinstalltorch2.1.0torchvision0.16.0torchaudio2.1.0 --index-url https://download.pytorch.org/whl/cpu⚠️踩坑预警装的是 CPU 版本的 PyTorch不是 CUDA 版本。pytorch-adapter 会替换 PyTorch 的后端从 CUDA 换成 NPU如果你装了 CUDA 版本会冲突。第3步装 pytorch-adapter# 方法1pip 安装推荐pipinstallpytorch-npu3.0.0# 对应 CANN 8.0# 方法2源码编译如果你想改 adapter 的代码gitclone https://atomgit.com/cann/pytorch-adapter.gitcdpytorch-adaptergitcheckout 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.fc1nn.Linear(784,512)self.fc2nn.Linear(512,10)defforward(self,x):xtorch.relu(self.fc1(x))xself.fc2(x)returnx modelMyModel()# 2. 把模型搬到 NPU 上关键modelmodel.npu()# 类似 .cuda()# 3. 把输入数据搬到 NPU 上input_datatorch.randn(32,784).npu()# 类似 .cuda()# 4. 前向计算自动在 NPU 上跑outputmodel(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. 定义模型 搬到 NPUmodelMyModel().npu()# 2. 定义损失函数 优化器跟 GPU 训练一模一样criterionnn.CrossEntropyLoss()optimizeroptim.Adam(model.parameters(),lr0.001)# 3. 训练循环跟 GPU 训练一模一样forepochinrange(10):forbatch_idx,(data,target)inenumerate(train_loader):# 把数据搬到 NPUdata,targetdata.npu(),target.npu()# 前向计算outputmodel(data)losscriterion(output,target)# 反向传播自动在 NPU 上算梯度optimizer.zero_grad()loss.backward()optimizer.step()ifbatch_idx%1000:print(fEpoch{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 上modelMyModel().npu()# 先搬到 NPUmodel.load_state_dict(torch.load(model_npu.pth))model.eval()# 推理模式# 3. 推理input_datatorch.randn(1,784).npu()outputmodel(input_data)predoutput.argmax(dim1)print(fPrediction:{pred.item()})⚠️踩坑预警保存的模型是 NPU 格式的不是 CPU 格式。如果你要拿到 CPU 上推理得先搬到 CPUmodel model.cpu()再保存。性能调优让模型跑得更快调优点1开启算子融合Graph Modepytorch-adapter 支持 Graph Mode图模式会自动融合算子类似 GE 图引擎的功能性能提升 20%-30%。importtorch# 开启 Graph Mode在模型推理前调用torch.npu.enable_graph_mode(True)# 推理会自动做算子融合input_datatorch.randn(1,784).npu()outputmodel(input_data)⚠️踩坑预警Graph Mode 不是所有模型都支持。如果你的模型有动态控制流if/for/whileGraph Mode 会报错。这时候得关掉 Graph Modetorch.npu.enable_graph_mode(False)。调优点2用 NPU 的加速库ATB如果你用的是 Transformer 类模型BERT/GPT/LLaMA可以用 ATBAscend Transformer Boost加速性能提升 50%-100%。fromatbimportTransformerInfer# 用 ATB 包装模型modelMyGPTModel().npu()inferTransformerInfer(model,max_seq_len2048,soc_typeAscend910)# 推理自动用 ATB 加速input_idstorch.tensor([[1,2,3,4,5]]).npu()outputinfer.forward(input_ids)调优点3用混合精度训练FP16NPU 的 FP16 算力是 FP32 的 2-3 倍。用混合精度训练性能提升 50%-80%。importtorch# 开启自动混合精度跟 GPU 的 amp 用法一样withtorch.npu.amp.autocast():outputmodel(input_data)losscriterion(output,target)loss.backward()踩坑实录踩坑1模型搬不到 NPU 上.npu() 报错原因PyTorch 版本跟 pytorch-adapter 不匹配。解决方案检查版本对应关系PyTorch 版本pytorch-adapter 版本1.13v2.0.x2.0v3.0.x2.1v3.1.x踩坑2训练时梯度算不出来loss.backward() 报错原因用了自定义的 autograd 函数torch.autograd.Functionadapter 没法自动替换。解决方案手动改自定义函数的代码把ctx.cuda()改成ctx.npu()。踩坑3推理性能不达标比 GPU 慢原因没开 Graph Mode 或者没用 ATB 加速库。解决方案开启 Graph Modetorch.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-adapterhttps://atomgit.com/cann/ascend-transformer-boosthttps://atomgit.com/cann/runtime