小模型在昇腾NPU上的推理过去被很多人当成挂个框架就能跑的事实际操作下来并不是那么无脑。我在用昇腾910B系列的服务器部署文本分类、embedding、reranker这类小模型时走过不少弯路尤其是torch_npu的迁移细节网上教程点到为止的多能讲透的少。这篇就把我自己从环境准备、torch_npu模型迁移到最终推理部署的完整过程整理出来希望对准备在昇腾NPU上做推理的同行有点参考价值。文章主要面向AI应用工程师、算法工程师和推理平台维护者也适合想了解昇腾NPU部署的新手。我会用一个中文BERT文本分类模型做主线案例把踩过的坑和排查思路一并分享。看完你就能理解小模型在NPU上做推理部署为什么值得专门走一遍torch_npu迁移。1. 为什么小模型在NPU上要专门做torch_npu迁移1.1 CPU、GPU、NPU部署选型到底怎么判断芯片领域里CPU、GPU、TPU、NPU这些概念初看确实让人头大。我见过很多纯粹做算法的人对硬件选型判断比较简单训练用GPU推理也顺手用GPU。这种做法不是不行但局限性很明显。CPU是通用的控制核心擅长复杂逻辑和分支跳转计算密度不高GPU最初为图形渲染设计通用矩阵计算和并行能力都很强适合训练和通用AI任务TPU是谷歌的张量处理单元主攻深度学习的专用场景NPU则是专门面向神经网络计算的处理器昇腾就是这一类的代表。昇腾NPU的AI Core内部有主网格阵列main grid array和MAC阵列针对矩阵乘加运算做了专门优化。这个架构让NPU在相同功耗下跑神经网络推理理论上能比通用处理器更高效。真实部署环境里成本、功耗、配套生态是三个关键因素。小模型的推理场景——比如embedding向量、reranker精排、文本分类、小参数量的对话模型——计算量不大但调用频次很高部署在CPU上延迟未必能接受部署在GPU上又有些浪费昇腾NPU这样的专用AI芯片反而显得更合适。尤其在实际项目中昇腾服务器已经采购到位的情况下能不能用好这块算力直接决定项目交付的成本和效率。1.2 为什么不能直接把PyTorch模型拿来跑NPUPyTorch官方目前只原生支持英伟达GPU和Apple Silicon上的MPS后端。想在非英伟达硬件上用PyTorch通常得靠厂商提供的加速插件。昇腾的PyTorch适配层就是torch_npu它是一个动态库插件作用是在PyTorch底层把算子调度交到昇腾的AI Core上执行。你写的仍然是熟悉的PyTorch代码但实际计算已经发生在NPU上这就是torch_npu的核心价值。换句话说torch_npu解决的就是模型迁移这件事。没有它你的model.cuda()会直接报错因为NPU不是cuda设备有了它你只需要做很少量的代码调整就可以把在GPU上验证过的模型跑在NPU上。这里要强调一个容易被忽略的点torch_npu不是把PyTorch翻译成另一套框架而是让PyTorch认识NPU这个设备。所以原有模型代码、训练逻辑、推理逻辑大部分不需要动改的是设备相关的胶水代码。1.3 小模型场景下torch_npu比大推理框架更实用现在提到推理框架很多人第一反应是vLLM。vLLM在LLM大模型场景下确实成熟但在昇腾NPU上用vLLM跑embedding和reranker这类模型实际过程里就会感受到现有编排对这种场景的支持并不理想甚至部分模型无法通过vLLM启动。原因在于vLLM的目标是大文本生成模型它的调度核心围绕KV Cache和连续batching而embedding、reranker这些模型不需要KV Cache也没有batch内解码过程强行套框架反而画蛇添足。所以我的建议是小模型推理优先考虑torch_npu手写推理服务。这样有几个好处依赖简单只要torch_npu和transformers控制力强每个函数自己掌控调试不黑盒部署轻量不带几百MB的框架依赖扩展灵活想接HTTP服务、消息队列都方便。这篇里我就用torch_npu直接写推理你会发现这条路完全可行。2. 环境准备昇腾NPU软件栈安装与自检先把话说在前面昇腾环境里版本匹配就是最大的坑。我见过太多人模型代码没问题最后卡在驱动、CANN和torch_npu版本不匹配上跑起来各种诡异报错。所以这部分内容我建议你跟着仔细过一遍别跳步。2.1 硬件自检用npu-smi确认设备状态拿到服务器第一件事不是急着装环境而是先确认NPU设备状态。昇腾服务器一般会带一个npu-smi命令类似GPU场景下的nvidia-smi。在终端执行npu-smi info正常能看到卡列表、算力利用率、显存占用等信息。如果这个命令能查出来但你的进程依然无法用NPU大概率是驱动权限或者容器映射问题。如果是root用户检查一下是否有权限隔离如果是容器部署需要确认启动容器时是否把设备节点正确映射进来了。另一个常用环境变量是ASCEND_RT_VISIBLE_DEVICES它的作用是控制当前进程可见的NPU设备和CUDA_VISIBLE_DEVICES非常像。多卡服务器上只使用某张卡就在运行前设置export ASCEND_RT_VISIBLE_DEVICES0这一步很多人容易漏。如果一台机器上有8张NPU卡你没设置这个变量进程可能默认申请全部设备的资源最终因为其他设备被别的任务占用而出错。2.2 驱动、固件和CANN工具包的版本匹配昇腾NPU的计算体系分三层最下面是固件和驱动负责硬件管理和资源分配中间是CANN提供算子库、图编译引擎和运行时最上层是PyTorch加torch_npu这样的框架适配。所以安装顺序不能乱先装固件和驱动再装CANN最后装torch_npu。CANN里有几个关键组件需要了解ACL运行时库负责设备管理AscendCL应用开发接口提供上层API还有专用的昇腾BLAS库做矩阵运算加速。很多人没意识到CANN自带针对昇腾NPU优化的BLAS实现PyTorch在NPU上的矩阵运算最终会走到这套底库上这是NPU性能能发挥出来的关键之一。版本匹配的原则很简单驱动和固件版本必须能支持你选定的CANN版本CANN版本必须能支持你选定的torch_npu版本torch_npu又和PyTorch版本一一对应。这三层任何一个错位都可能出现环境看起来装好了一跑就崩的局面。这里有一个非常重要的经验不要装最新版本的torch_npu而是查清楚你的CANN版本对应哪个torch_npu版本用官方版本表确认关系。组件版本选择建议说明固件与驱动按服务器型号从官方支持矩阵选参考CANN版本对驱动的最低要求CANN Toolkit与驱动兼容的版本优先用昇腾官方推荐组合Python3.8及以上需和torch/torch_npu编译目标一致PyTorch与torch_npu官方映射表对应比如2.1.0系列对应特定torch_npu版本torch_npu和torch版本精确对应版本号后常有post字样安装时看清楚2.3 torch_npu安装pip安装与源码编译的选择理论上torch_npu可以源码编译但一般情况下用pip直接安装是最省事的。昇腾官方会发布和PyTorch版本对应的预编译包安装命令形如pip install torch2.1.0 pip install torch_npu2.1.0.post6注意这里的torch_npu版本号不是随意的尾部post后面的数字对应不同的CANN版本。我曾经因为图省事直接装了不匹配的版本结果运行时提示算子库初始化失败。因此安装前先确认好CANN版本再去找适配的torch_npu版本。源码编译适合有特殊定制需求的场景一般不建议。编译需要下载CANN源码包组件操作成本高而且编译时间不短。如果你只是做推理部署pip预编译包足够用了。装完之后还有环境变量的配置这是新手最容易漏的环节。典型配置如下source /usr/local/Ascend/ascend-toolkit/set_env.sh export LD_LIBRARY_PATH/usr/local/Ascend/ascend-toolkit/latest/lib64:$LD_LIBRARY_PATH如果环境变量没配好最典型的报错是运行时找不到libascendcl.so之类的动态库。这类错误和模型代码没关系纯环境问题排查起来只要对照环境变量配置就能定位。2.4 用两段代码快速验证环境装完环境后不要急着跑大模型先用最小验证代码确认NPU能正常起算子。第一段验证设备可用import torch import torch_npu print(torch_npu.npu.is_available()) print(torch_npu.npu.device_count())如果输出True和大于0的数字说明基础设备OK。第二段验证真正的计算链路import torch import torch_npu a torch.randn(4, 4).npu() b torch.randn(4, 4).npu() c torch.matmul(a, b) print(c.cpu())这里刻意在NPU上做一次矩阵乘目的是确认算子库链路完好。如果这一步跑通你的torch_npu迁移环境就算准备好了。我第一次跑这段的时候遇到了libascendcl.so找不到的问题就是因为环境变量没配对source了set_env.sh之后就好了。3. torch_npu模型迁移的核心步骤与原理这部分是全文的干货重心我会把迁移涉及到的所有关键代码和背后逻辑讲清楚你照着改就行。3.1 四步迁移法从CUDA习惯切到NPU习惯假设你原来有一段PyTorch推理代码模型是model输入是input_ids。如果代码是用GPU写的通常会有model.cuda()、input_ids.cuda()这样的调用。迁移到NPU标准操作就四步。第一步导入torch_npuimport torch import torch_npu这一步的作用是注册NPU后端。不import后面的torch.npu调用根本不存在。第二步设置设备device torch.device(npu:0)如果你使用多卡npu:0、npu:1这样的写法含义和cuda:0是一致的。第三步把模型搬到NPUmodel.to(device)第四步把输入数据搬到NPUinput_ids input_ids.to(device) attention_mask attention_mask.to(device)就这么简单。很多模型代码其实只需要改设备相关的地方剩下的前向逻辑、后处理逻辑完全可以不动。这也是torch_npu设计的精髓——它尽量让你少改代码。这里我要专门说一个我犯过的低级错误输入数据的device和模型的device不一致。这个问题在CPU/GPU混合的时候也常见但NPU更隐蔽因为报错有时不会直接提示device mismatch而是会在某个算子执行时报奇怪的shape错误。排查思路很简单在推理前打印一下各对象的device属性print(model.device) print(input_ids.device)一旦发现一个是npu:0一个是cpu问题就清楚了。3.2 推理部署的细节别把训练习惯带进来迁移到位后真正写推理部署代码时有几个关键点容易被忽略。第一个是model.eval()。这个调用会关闭dropout和batch normalization的统计更新。很多人在GPU推理的时候都会记得调用但迁移到NPU上后偶尔会因为想快速验证忘了置成eval模式导致结果不稳定。第二个是torch.no_grad()。推理阶段不计算梯度可以减少显存占用和计算开销这个无论CPU、GPU还是NPU都一样重要。第三个是固定输入长度。NPU和GPU在动态shape处理上有个显著差异GPU对动态shape的容忍度相对高但NPU某些算子在动态shape下会重新走一遍图编译导致推理延迟变大。实际操作时我建议把输入统一padding到固定长度比如128或256并关闭tokenizer的动态paddinginputs tokenizer( texts, max_length128, paddingmax_length, truncationTrue, return_tensorspt, )第四个是npu同步。推理结束后如果想统计耗时不要只看Python层面经过的时间因为NPU的算子执行是异步的可能代码已经走完但算子还在设备上跑。这时候需要调用torch.npu.synchronize()强制同步保证耗时统计准确这也是做性能评估时的常规操作。3.3 混合精度与图编译小模型性能优化的两个抓手小模型在NPU上默认跑FP32可以工作但如果你想进一步压榨性能可以考虑混合精度。PyTorch的自动混合精度机制沿用到了torch_npu推理时把部分算子用FP16执行对速度有明显提升。文本分类这类任务对精度敏感度一般不高FP16基本无忧。混合精度的代码写法以我的使用经验来看可以直接用with torch.autocast(device_typenpu, dtypetorch.float16): outputs model(**inputs)如果你的torch_npu版本较老有的版本也支持torch.npu.amp.autocast()写法以实际安装版本的文档为准。另一个优化方向是图模式。torch_npu支持将PyTorch动态图模式转换到静态图编译减少算子调度开销。开启方式大概如下不同版本名称略有差异import torch_npu torch_npu.npu.set_compile_mode(jit_compileTrue)开启图模式后首次执行时会有编译开销第二次开始速度才有明显提升。但需要注意不是所有算子都能完美走静态图如果遇到编译失败可以先关掉图模式先跑通功能再调优。我自己的经验是小模型里文本编码模型的算子比较规整图模式收益稳而一些涉及复杂控制流的自定义模型收益反而不稳定。3.4 从CUDA代码迁移时的机械替换清单如果你手头的模型代码还不干净直接用cuda写法迁移时记住这张替换表原CUDA写法NPU写法说明model.cuda()model.to(npu:0)统一用to方式更清晰input.cuda()input.to(npu:0)同上torch.cuda.is_available()torch_npu.npu.is_available()设备判断torch.cuda.device_count()torch_npu.npu.device_count()设备数量device cudadevice npu:0设备字符串torch.cuda.synchronize()torch.npu.synchronize()同步等待这表看着简单但逐行执行时容易漏。有些代码里会出现硬编码的cuda字符串直接改成npu即可。还有加载模型权重时的map_location参数如果从GPU训练好的权重加载到NPU跑推理要写成map_locationnpu:0。这一条我实际遇到过权重能加载但设备对不上最后排查了半天才发现是map_location的问题。4. 完整案例BERT中文分类模型在昇腾NPU上的推理部署前面讲完思路这里放一个完整可跑的案例。我用的是BERT-base-chinese文本分类模型这个模型参数规模约110M属于典型的小模型适合做演示。完整代码可以在你的昇腾环境上直接改路径后运行。4.1 模型与依赖准备我的环境假设是PyTorch 2.1.x加torch_npu 2.1.0.post系列加transformers 4.x加Python 3.9。transformers用于加载模型和tokenizer。模型文件从官方仓库拉下来或者在离线环境里手动拷贝到本地目录。pip install transformers这里有一个实际部署的经验昇腾NPU服务器很多是内网环境无法直接从huggingface下载模型。建议在能联网的机器上把模型和tokenizer预下载好然后整体拷贝到NPU服务器的本地目录。加载时直接用本地目录路径model_path /data/models/bert-base-chinese tokenizer AutoTokenizer.from_pretrained(model_path) model AutoModelForSequenceClassification.from_pretrained(model_path, num_labels2)4.2 完整推理脚本从加载到输出一应俱全我的案例是二分类情感判断输入一句中文文本输出正面或负面的判断。整体脚本结构如下import time import torch import torch_npu from transformers import AutoTokenizer, AutoModelForSequenceClassification # 1. 设置NPU设备 device torch.device(npu:0) # 2. 加载模型和tokenizer model_path /data/models/bert-base-chinese tokenizer AutoTokenizer.from_pretrained(model_path) model AutoModelForSequenceClassification.from_pretrained(model_path, num_labels2) model.to(device) model.eval() # 3. 定义推理函数 def predict(text): inputs tokenizer( text, max_length128, paddingmax_length, truncationTrue, return_tensorspt, ) inputs {k: v.to(device) for k, v in inputs.items()} with torch.no_grad(): outputs model(**inputs) logits outputs.logits pred torch.argmax(logits, dim-1) return int(pred.cpu()) # 4. 预热NPU首次推理有算子编译开销先跑一次 _ predict(这个项目交付质量不错) torch.npu.synchronize() # 5. 性能测试连续推理多次统计平均耗时 texts [ 这个项目交付质量不错我非常满意, 今天遇到的问题太多了体验很差, 速度虽然一般但整体功能完整, ] for text in texts: start time.time() result predict(text) torch.npu.synchronize() cost_ms (time.time() - start) * 1000 print(f文本: {text} | 预测: {result} | 耗时: {cost_ms:.2f}ms)上面这段代码里有几个点专门说明一下。predict函数内部的tokenizer输出是Python字典我通过字典推导式把它搬到NPU这比逐个赋值更干净。返回结果时用pred.cpu()把张量从NPU搬回CPU因为后面print、json序列化都要在CPU上操作。整个过程就是一个标准的torch_npu推理链路。预热步骤千万别删。我在实际使用中测过不预热的情况下首次推理可能要几百毫秒到一秒以上预热后单条推理能稳定下降到几十毫秒甚至更高性能。这不是模型问题是NPU算子在首次执行时需要做编译和资源分配属于正常现象。4.3 看性能同样模型CPU和NPU的差距我在同一台服务器上简单对比过这个模型CPU推理单条耗时大约在150到300毫秒波动取决于CPU负载和线程数NPU预热后单条耗时能压到几十毫秒提升还是很明显的。如果你开启混合精度和图模式还会更低。这个量级的小模型单看绝对耗时不算大但在高并发场景下优势会积累得很明显。这里要提醒一句做性能测试一定要等设备空闲且要多测几轮取平均值不要用首次推理的时间作为结论否则很容易得出NPU比CPU还慢这种错误结论。我第一次测性能时就是没预热直接拿着第一次的耗时去对比差点被数据误导。4.4 把推理函数封装成HTTP服务脚本能跑通后部署阶段通常需要把推理能力暴露成HTTP接口。我用FastAPI封装整体结构如下from fastapi import FastAPI, Request import uvicorn app FastAPI() app.post(/predict) async def predict_api(request: Request): data await request.json() text data.get(text, ) result predict(text) return {text: text, label: result} if __name__ __main__: uvicorn.run(app, host0.0.0.0, port8080, workers1)注意workers要设为1这是很多新手踩坑的地方。模型加载后常驻内存如果开多个worker每个进程都会完整加载一份模型内存占用成倍增加而且多个进程同时访问NPU设备如果没有合理调度反而会互相干扰。小模型的QPS需求一般单进程就能扛住真需要水平扩展时再考虑多进程部署和负载均衡。另外FastAPI默认的线程池和torch推理的线程设置有时候会打架。如果你发现并发请求响应不稳定可以设置一下CPU线程数torch.set_num_threads(1)这样能减少推理时的线程争抢实测对稳定性有帮助。5. 常见问题与排查技巧实录最后把我在实际部署过程中遇到的高频问题整理成速查表然后逐个展开讲排查思路全是实战中沉淀下来的一手经验。5.1 NPU设备不可用从环境变量到固件驱动这个问题表现很多样torch_npu.npu.is_available()返回False或者创建context报错甚至import阶段就直接崩。排查顺序如下。第一步