翻译模型的部署一直是件烦心事,尤其是把 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_ENDPOINT=https://hf-mirror.com然后正常调用from_pretrained即可,完全不用改代码逻辑。完整下载流程我建议这样写:
from huggingface_hub import snapshot_download snapshot_download( repo_id="Helsinki-NLP/opus-mt-en-zh", local_dir="./local_hf_model", )这样可以一次拉取全部权重和 tokenizer 文件,避免后续导出时还要反复访问网络。下载完以后检查一下目录,确认pytorch_model.bin、config.json、vocab.json、spiece.model这些关键文件都在。缺文件是后面报错的高频原因。
3. 核心导出过程:拆开 encoder 与 decoder
3.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_dtype=torch.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_ids=input_ids, attention_mask=attention_mask, return_dict=False, ) 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), dtype=torch.long) dummy_mask = torch.ones(1, 16, dtype=torch.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_version=14, do_constant_folding=True, )这里最关键的是dynamic_axes。翻译请求的句子长度不可能固定,所以batch和src_len两个维度都必须标成动态。如果不标,导出后模型只能接受固定长度的输入,这在真实场景里完全没法用。
3.3 导出 decoder
decoder 比 encoder 复杂一点,因为它后面还接了一个 lm_head 输出词表 logits:
class 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_ids=decoder_input_ids, attention_mask=decoder_attention_mask, encoder_hidden_states=encoder_hidden_states, return_dict=False, ) 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), dtype=torch.long) dummy_decoder_mask = torch.ones(1, 16, dtype=torch.long) dummy_enc_hidden = torch.randn(1, 16, model.config.d_model, dtype=torch.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_version=14, do_constant_folding=True, )这里有个容易搞混的点: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_input="decoder_sim.onnx", model_output="decoder_int8.onnx", weight_type=QuantType.QInt8, per_channel=True, )为什么只量化 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_options=session_options, providers=["CPUExecutionProvider"], ) dec_session = ort.InferenceSession( "decoder_int8.onnx", sess_options=session_options, providers=["CPUExecutionProvider"], )graph_optimization_level拉到ORT_ENABLE_ALL,ONNX Runtime 会尽量做算符融合和内存规划。intra_op_num_threads可以根据机器核数调,设太高反而会带来线程切换开销。
5.3 贪心解码主循环
def translate(text, max_length=128): # 编码源语言 encoded = tokenizer(text, return_tensors="pt", truncation=True, max_length=256) 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]], dtype=np.int64) for _ in range(max_length): dec_mask = np.ones_like(decoder_ids, dtype=np.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]], dtype=np.int64)], axis=-1, ) return tokenizer.decode(decoder_ids[0].tolist(), skip_special_tokens=True)有几个细节需要说明。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 生成方案写出来,如果你现在照着本文跑通了,不妨多试试不同模型,踩过的坑都会变成你自己的经验。