ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

ONNX Runtime CPU部署英译中模型:从Helsinki迁移到INT8量化实战

ONNX Runtime CPU部署英译中模型:从Helsinki迁移到INT8量化实战 翻译模型的部署一直是件烦心事尤其是把 HuggingFace 上的英译中模型搬到生产环境。你在 PyTorch 里跑得挺好一上服务就发现 CPU 吃满、响应超时。与其折腾各种加速编译我建议先把模型迁到 ONNX再把推理交给 ONNX Runtime。最近我正好把Helsinki-NLP/opus-mt-en-zh这个经典的英译中模型完整迁移了一遍顺手做了 INT8 量化前前后后踩了不少坑这篇就当是给后来人留的一份操作记录。这个流程适合两类人一是准备上线翻译服务、接受不起 PyTorch 动态图开销的后端开发二是想做离线翻译工具、希望模型体积更小的客户端工程师。整条链路从下载权重到 ONNX 推理大约半小时能跑通。我不会只贴能跑通的代码还会把每一步背后的取舍讲清楚。1. 为什么要做这个迁移1.1 线上翻译服务的成本压力翻译模型的响应延迟和并发量是一对天然矛盾。用 PyTorch 直接跑 MarianMT 这种 300MB 级别的 encoder-decoder 模型CPU 上单条短句推理通常要 200 到 400 毫秒一旦并发上来延迟直接翻倍。问题的根源不只是模型大更在于 PyTorch 动态图机制在每步运算里都要维护自动求导图这在离线训练时是优势在在线推理时就变成了纯开销。ONNX Runtime 则完全换了一套思路。它先把模型计算图做静态优化再根据不同的 CPU 指令集做算子融合。Transformer 里大量连续的 MatMul、Add、LayerNorm 会被合并成少数几个专用算子省掉中间张量的读写。我实测下来同一个模型在相同硬件上ONNX Runtime 的 CPU 推理耗时普遍能比 PyTorch eager 模式快 30% 到 50%输入越长收益越明显。1.2 部署形态的灵活性模型导出成 ONNX 文件后部署形态就打开了。ONNX 是开放交换格式ONNX Runtime、OpenVINO、TensorRT、NCNN 都能加载。也就是说同一份模型文件既能在服务器 CPU 上跑也能在 GPU 上跑还能搬到边缘设备上。换成 PyTorch 原生的 TorchScript 就没这么自由绑得太死。更重要的是运行时依赖变轻了。PyTorch 推理服务动不动要求几百 MB 的 torch 库还要处理和 CUDA、MKL 的版本兼容问题。ONNX 这边基本一个onnxruntime包就能搞定对 Docker 镜像体积和启动速度都有明显改善。生产环境里少一个依赖就少一类事故。1.3 为 INT8 量化铺路量化是我这次迁移的第二动机。ONNX 生态里的动态量化工具已经非常成熟quantize_dynamic一条命令就能把模型权重从 FP32 压到 INT8而且不需要重新训练。MarianMT 这种翻译模型decoder 在自回归生成时会被反复执行量化它对内存带宽和计算延迟的改善都相当可观。2. 模型选型与前期准备2.1 用一个轻量的英译中模型做主线HuggingFace 上有不少英译中模型为什么选Helsinki-NLP/opus-mt-en-zh因为它是经典的 MarianMT 架构encoder-decoder 结构非常规整参数规模约 300MB中文效果虽然比不上 NLLB 那种大模型但对一个部署示范项目来说足够典型。NLLB 和 M2M100 效果当然更好但它们的导出过程涉及更多特殊处理新手很容易卡在算子的兼容性上。模型选型时要想清楚一个前提你是在做效果验证还是在做生产部署。如果是验证全链路可行性选结构简单的 MarianMT 能把精力都放在 ONNX 流程上如果确实需要高翻译质量也建议先按本文流程跑通再换成目标模型导出逻辑基本复用。2.2 环境与依赖安装我这边的环境是 Python 3.10、PyTorch 2.1核心依赖如下pip install torch transformers onnx onnxruntime onnxslimonnxslim用来做计算图简化比老牌的 onnx-simplifier 更省心。optimum工具包我也装了虽然最后没有用它做主力导出但可以拿来作对比验证。如果是 Windows 环境注意onnxruntime默认走 CPU想用 GPU 需要单独装onnxruntime-gpu不过本文场景 CPU 就足够了。2.3 把模型下载到本地很多人会卡在这一步。HuggingFace 官网在国内的连接质量时好时坏我的办法是提前把环境变量指向镜像站export HF_ENDPOINThttps://hf-mirror.com然后正常调用from_pretrained即可完全不用改代码逻辑。完整下载流程我建议这样写from huggingface_hub import snapshot_download snapshot_download( repo_idHelsinki-NLP/opus-mt-en-zh, local_dir./local_hf_model, )这样可以一次拉取全部权重和 tokenizer 文件避免后续导出时还要反复访问网络。下载完以后检查一下目录确认pytorch_model.bin、config.json、vocab.json、spiece.model这些关键文件都在。缺文件是后面报错的高频原因。3. 核心导出过程拆开 encoder 与 decoder3.1 为什么不用 optimum 一键导出HuggingFace 官方工具链支持一条命令导出optimum-cli export onnx --model Helsinki-NLP/opus-mt-en-zh onnx/它能生成一个完整的 ONNX 模型输入是input_ids和attention_mask输出是 logits。但这里有个坑这个导出方式把整个 encoder decoder 塞进了一张计算图里。你每生成一个 token整棵树都要从头跑一遍encoder 的结果也反复重算自回归解码的效率完全发挥不出来。真实部署里我们希望源语言只编码一次把 encoder 的 hidden states 存下来之后每一步只跑 decoder。所以正确做法是手动导出两个 ONNX 文件一个 encoder、一个 decoder。这样生成的复杂度才符合预期。3.2 导出 encoder先加载模型并保持 eval 状态import torch from transformers import MarianMTModel, MarianTokenizer model_id Helsinki-NLP/opus-mt-en-zh tokenizer MarianTokenizer.from_pretrained(model_id) model MarianMTModel.from_pretrained(model_id, torch_dtypetorch.float32) model.eval()封装一个只暴露 encoder 的模块class EncoderWrapper(torch.nn.Module): def __init__(self, model): super().__init__() self.model model def forward(self, input_ids, attention_mask): outputs self.model.model.encoder( input_idsinput_ids, attention_maskattention_mask, return_dictFalse, ) return outputs[0]注意这里用的是model.model.encoder因为MarianMTModel内部还有一个MarianModel。不同版本 transformers 内部结构可能会微调导出前最好先打印一下模型结构确认路径。然后构造 dummy 输入并导出encoder EncoderWrapper(model) encoder.eval() dummy_input_ids torch.randint(0, tokenizer.vocab_size, (1, 16), dtypetorch.long) dummy_mask torch.ones(1, 16, dtypetorch.long) torch.onnx.export( encoder, (dummy_input_ids, dummy_mask), encoder.onnx, input_names[input_ids, attention_mask], output_names[encoder_hidden_states], dynamic_axes{ input_ids: {0: batch, 1: src_len}, attention_mask: {0: batch, 1: src_len}, encoder_hidden_states: {0: batch, 1: src_len}, }, opset_version14, do_constant_foldingTrue, )这里最关键的是dynamic_axes。翻译请求的句子长度不可能固定所以batch和src_len两个维度都必须标成动态。如果不标导出后模型只能接受固定长度的输入这在真实场景里完全没法用。3.3 导出 decoderdecoder 比 encoder 复杂一点因为它后面还接了一个 lm_head 输出词表 logitsclass DecoderWrapper(torch.nn.Module): def __init__(self, model): super().__init__() self.model model def forward(self, decoder_input_ids, decoder_attention_mask, encoder_hidden_states): outputs self.model.model.decoder( input_idsdecoder_input_ids, attention_maskdecoder_attention_mask, encoder_hidden_statesencoder_hidden_states, return_dictFalse, ) logits self.model.lm_head(outputs[0]) return logits导出decoder DecoderWrapper(model) decoder.eval() dummy_decoder_ids torch.randint(0, tokenizer.vocab_size, (1, 16), dtypetorch.long) dummy_decoder_mask torch.ones(1, 16, dtypetorch.long) dummy_enc_hidden torch.randn(1, 16, model.config.d_model, dtypetorch.float32) torch.onnx.export( decoder, (dummy_decoder_ids, dummy_decoder_mask, dummy_enc_hidden), decoder.onnx, input_names[decoder_input_ids, decoder_attention_mask, encoder_hidden_states], output_names[logits], dynamic_axes{ decoder_input_ids: {0: batch, 1: dec_len}, decoder_attention_mask: {0: batch, 1: dec_len}, encoder_hidden_states: {0: batch, 1: src_len}, logits: {0: batch, 1: dec_len}, }, opset_version14, do_constant_foldingTrue, )这里有个容易搞混的点encoder_hidden_states的第二个维度是src_len而decoder_input_ids的第二个维度是dec_len。它们各自独立变化所以在dynamic_axes里分别用不同的维度名称否则导出时会报符号冲突。3.4 验证导出的模型有没有跑偏导出只是第一步验证一致性更不能省。简单写一段脚本分别用 PyTorch 和 ONNX Runtime 跑同一个 dummy 输入比较输出import numpy as np import onnxruntime as ort encoder_session ort.InferenceSession(encoder.onnx, providers[CPUExecutionProvider]) decoder_session ort.InferenceSession(decoder.onnx, providers[CPUExecutionProvider]) with torch.no_grad(): torch_enc_out encoder(dummy_input_ids, dummy_mask).numpy() torch_dec_out decoder(dummy_decoder_ids, dummy_decoder_mask, torch.from_numpy(torch_enc_out)).numpy() onnx_enc_out encoder_session.run(None, { input_ids: dummy_input_ids.numpy(), attention_mask: dummy_mask.numpy(), })[0] onnx_dec_out decoder_session.run(None, { decoder_input_ids: dummy_decoder_ids.numpy(), decoder_attention_mask: dummy_decoder_mask.numpy(), encoder_hidden_states: onnx_enc_out, })[0] print(encoder cos sim:, np.corrcoef(torch_enc_out.ravel(), onnx_enc_out.ravel())[0, 1]) print(decoder cos sim:, np.corrcoef(torch_dec_out.ravel(), onnx_dec_out.ravel())[0, 1])浮点计算有误差很正常余弦相似度在 0.99 以上就算通过。如果低于这个值优先检查模型是否被不小心设成了 train 模式或者 ONNX 导出时的do_constant_folding是不是出了问题。4. 模型轻量化与 INT8 量化4.1 用 onnxslim 把计算图清理干净PyTorch 导出 ONNX 时计算图里会有不少冗余算子比如连续的 Reshape、Transpose、Identity。这些不影响正确性但会增加推理时的调度开销。用onnxslim一行就能收拾干净onnxslim encoder.onnx encoder_sim.onnx onnxslim decoder.onnx decoder_sim.onnx简化过后文件大小可能略有下降更重要的是图结构更利于 ONNX Runtime 做算子融合。我建议量化之前先做这一步否则某些中间节点的特殊形状可能干扰量化工具。4.2 对 decoder 做 INT8 动态量化量化工具直接用 ONNX Runtime 自带的from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( model_inputdecoder_sim.onnx, model_outputdecoder_int8.onnx, weight_typeQuantType.QInt8, per_channelTrue, )为什么只量化 decoder不量化 encoder因为自回归生成时 encoder 只跑一次decoder 要跑 N 步N 是输出序列长度通常几十个 token。decoder 的计算总量和权重访存压力都大得多量化它收益最明显。encoder 保持 FP32 反而能保留源语言编码精度。quantize_dynamic属于动态量化它不是把整个计算图都压成 INT8而是把权重从 FP32 映射到 INT8推理时在节点内反量化为 FP32 后再计算。这种方式对精度影响相对小同时能显著减少内存占用和访存带宽压力。4.3 量化的实际收益在我这边的参考机上decoder_sim.onnx 约 120MB量化之后只有 35MB 左右。CPU 上短句翻译耗时从 150ms 左右降到 90ms 左右长句收益更明显。当然这个数字受硬件和输入长度影响很大真正要紧的是量化的代价是否会影响翻译质量。我用几个常见句子做了对比量化前后的翻译基本一致。个别长难句会出现词汇选择的细微差异但没有出现乱码或者语序崩坏。如果发现量化后质量明显下滑可以考虑只量化 encoder 或者干脆放弃 INT8直接部署简化后的 FP32 模型这个取舍要在真实流量上做验证。5. 在 ONNX Runtime 里跑起完整翻译5.1 tokenizer 留在模型外ONNX 只替代神经网络部分tokenizer 还是需要用 transformers 加载。因为生产环境里不一定装了 transformers所以提前把 tokenizer 文件保存到本地目录运行时就从这个目录加载tokenizer.save_pretrained(./local_hf_model/tokenizer)后续部署时只需要带上这个 tokenizer 目录和两个 ONNX 文件完全不需要安装 torch。这是一个很容易被忽略的点很多人把精力都放在模型导出上忘了 tokenizer 其实也是一份需要单独交付的资产。5.2 创建 ONNX Runtime 会话导入必要的库设置会话import numpy as np import onnxruntime as ort session_options ort.SessionOptions() session_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL session_options.intra_op_num_threads 4 enc_session ort.InferenceSession( encoder_sim.onnx, sess_optionssession_options, providers[CPUExecutionProvider], ) dec_session ort.InferenceSession( decoder_int8.onnx, sess_optionssession_options, providers[CPUExecutionProvider], )graph_optimization_level拉到ORT_ENABLE_ALLONNX Runtime 会尽量做算符融合和内存规划。intra_op_num_threads可以根据机器核数调设太高反而会带来线程切换开销。5.3 贪心解码主循环def translate(text, max_length128): # 编码源语言 encoded tokenizer(text, return_tensorspt, truncationTrue, max_length256) src_ids encoded[input_ids].numpy().astype(np.int64) src_mask encoded[attention_mask].numpy().astype(np.int64) # encoder 只执行一次 enc_feed { input_ids: src_ids, attention_mask: src_mask, } enc_hidden enc_session.run(None, enc_feed)[0] # decoder 自回归生成 decoder_start model.config.decoder_start_token_id decoder_ids np.array([[decoder_start]], dtypenp.int64) for _ in range(max_length): dec_mask np.ones_like(decoder_ids, dtypenp.int64) dec_feed { decoder_input_ids: decoder_ids, decoder_attention_mask: dec_mask, encoder_hidden_states: enc_hidden, } logits dec_session.run(None, dec_feed)[0] next_token int(logits[0, -1, :].argmax()) if next_token tokenizer.eos_token_id: break decoder_ids np.concatenate( [decoder_ids, np.array([[next_token]], dtypenp.int64)], axis-1, ) return tokenizer.decode(decoder_ids[0].tolist(), skip_special_tokensTrue)有几个细节需要说明。decoder_attention_mask我直接传全 1因为 MarianMT 的 decoder 内部会自动构造因果掩码外部传的 mask 只是用来处理 padding这里没有 padding所以全部置 1 就能得到正确结果。另外这里导出的 decoder 没有包含 past_key_values 优化所以每一步都要把完整的decoder_ids喂进去序列越长计算量越大。想进一步提速就得单独导出带 KV Cache 的 decoder 版本那是另一个较复杂的工程以后可以单独开一篇。5.4 实测效果跑几个句子验证一下print(translate(The weather is nice today.)) print(translate(I would like to book a table for two.))输出效果基本符合预期语义和语序都比较通顺。ONNX Runtime 的推理结果和原模型极端接近差别只在浮点尾数。如果你的业务对某些专有名词特别敏感建议把量化前后的模型针对领域数据做一次系统评测再决定上线用哪个版本。6. 常见问题与排查实录6.1 导出时报 Unsupported operator现象torch.onnx.export中途直接抛异常说某个 ATen 算子不支持导出到 ONNX。原因大多是对应的 opset 版本太老部分算子没有在旧版本里定义。最常见的解决办法是把opset_version从 12 提到 14 或者 17。如果提高 opset 后仍然报错可以尝试把do_constant_folding改成False。有些算子是在常量折叠阶段被替换成 ONNX 不支持的格式关掉这个选项能让导出流程继续。注意这可能会让模型文件稍微变大但对推理正确性没有影响。6.2 变长输入维度冲突现象导出的模型用固定长度 dummy 输入测试没问题但正式运行时输入长度一变ONNX Runtime 就报维度不匹配。九成原因是dynamic_axes没有配全。比如说我在导出 decoder 时忘了把encoder_hidden_states的src_len标成动态结果源句子长度从 16 变成 20 时直接炸了。排查思路是看报错信息里是哪个节点的哪个维度不匹配然后回到torch.onnx.export的dynamic_axes里补全对应维度。凡是涉及 batch 和 seq 的所有输入输出都应该在dynamic_axes中声明一个都不能漏。6.3 量化后输出异常现象量化后的 decoder 跑出来结果全是一串无意义 token甚至出现 NaN。这种情况先别急着怀疑量化本身先检查量化前的模型是否正常。如果量化前正常再尝试缩小量化范围。一个比较实用的技巧是给quantize_dynamic加上op_types_to_quantize[MatMul]参数只量化最核心的矩阵乘算子避开 LayerNorm、Gather 这些容易出问题的节点。如果仍然不行可以退回 FP32使用简化后的模型部署。INT8 是优化项不是必选项服务稳定永远排在性能前面。6.4 模型文件下载不完整现象加载 tokenizer 时提示缺少spiece.model或者加载模型时找不到权重。这通常是因为从镜像站下载时网络中断snapshot_download只拉取了部分文件。解决方案是下载前检查本地目录或直接删除后重新执行snapshot_download拉取。下载完成建议做个简单校验确认pytorch_model.bin的大小和 HuggingFace 仓库页面上标注的文件大小一致。很多稀奇古怪的报错追溯到最后就是权重文件缺了几百 KB。最后再分享一点个人体会这套流程跑下来我最深刻的感受是模型导出不是最难的真正难的是搞清楚计算图里数据流向。手动拆 encoder 和 decoder 这个决定让推理性能和部署形态都有了质的提升。Helsinki-NLP/opus-mt-en-zh只是一个基本功练习同一套手法完全可以平移到 NLLB、T5 这些模型上甚至 YOLO、TTS、ASR 模型的 ONNX 部署思路也相通。下次我打算把带 KV Cache 的 decoder 导出和 beam search 生成方案写出来如果你现在照着本文跑通了不妨多试试不同模型踩过的坑都会变成你自己的经验。
返回列表