有个开发兄弟前两天找我帮忙他的需求很直接把 HuggingFace 模型仓库里现成的英译中模型迁到 ONNX跑在一台只有 CPU 的服务器上内存 8G还要接进现有的 Web 服务。翻译模型部署这件事难的不是把权重下载下来而是导出之后的 runtime 行为完全变了输入输出形状要自己管自回归循环要自己写连 tokenizer 都得单独保留。这篇东西就是我把从环境准备到输出对齐整条链路走完之后留下的实操记录里面有能直接抄的脚本也有我在这轮踩过的坑。适合两类人一类是正准备给翻译模型做 CPU 推理优化的开发另一类是想搞懂 seq2seq 模型在 ONNX 下怎么落地的算法工程师。1. 迁移思路拆解到底为什么迁又该在哪种场景下迁1.1 三个核心收益跑分、依赖、内存先说我为什么愿意趟这趟浑水。ONNX Runtime 对 Transformer 类模型的算子做了大量手工调优在 CPU 上跑同一套参数通常比裸 PyTorch 快 20% 到 50%具体取决于你的 op 融合情况和机器型号。这个加速不花你的钱只花一次导出的功夫。第二个收益是部署环境变轻。模型一旦导出成 ONNX目标服务器上就不需要再装一整套 torch 和 CUDA 依赖只要一个 onnxruntime 加 numpy 就能跑。别小看这一点很多内网服务器不允许随便装大体积 Python 包有的甚至要求离线安装。我用过一台只有 2G 磁盘配额的老机器torch 都不一定能塞进去但 onnxruntime 加相关依赖加起来不到 200M这就是迁移的现实意义。第三个收益是内存和显存占用。ONNX 静态图会让推理框架有机会复用中间缓冲区不像 PyTorch 每次 forward 都重新分配张量。在 8G 内存的服务器上模型加载后实际峰值能比原版 PyTorch 少 1 到 2G别觉得夸张真要跑高并发时这 2G 就是生死线。1.2 什么时候不该迁两个反向场景不是所有翻译模型都值得往 ONNX 上搬。如果业务还在频繁调整模型结构、改前处理逻辑、经常重新训练小版本那我劝你等等。ONNX 是静态图模型里每改一个算子就得重新导出、重新跑一致性验证这个成本叠加到迭代周期里很容易让人崩溃。另一个不适合迁移的场景是你的翻译流程深度绑定了 HuggingFace generate 里那些花哨采样方式。比如你需要 top-k 加 top-p 随机采样还要对多个候选结果做重排序那你用 ONNX 自绘解码循环的血泪程度会超出预期。ONNX 只负责前向计算采样策略、beam search、父序列回溯这些逻辑全部要自己写。所以如果业务只需要贪心解码或简单的 beam search迁移成本是可控的如果已经依赖复杂的自定义生成策略建议保持 PyTorch 推理。1.3 挑选英译中模型的三条标准HuggingFace 上叫“英译中”的模型不少但挑错模型会让你后面每一步都难受。我一般按三个标准筛。第一语言对必须准确。很多 MarianMT 系模型的名字和实际语言对并不完全对应一定要看模型卡里 source_language 和 target_language 的配置还要检查 tokenizer 词表里有没有目标语言的特殊符号。有的模型实际是英法换了个视觉相近的缩写粗心直接下载就会出大问题。第二模型体积要匹配你的部署资源。CPU 8G 内存就优先选 250M 参数以内的中等尺寸模型MarianMT 系列通常是个好选择如果业务必须用更大规模的多语言模型比如 M2M100 这类导出没问题但要提前算好 decoder 的 logits 层厚度输出词表动辄几万甚至十几万每次自回归都要算一次大矩阵乘CPU 下会非常肉疼。第三看架构导出成熟度。Bart、T5 这类 Encoder-Decoder 架构在 ONNX 生态里支持最完整Marian 架构和它们类似通过 optimum 也能顺利导出但不少多语言模型的 tokenizer 有自定义逻辑导出时有意想不到的分支你需要有手改封装的能力。综合下来第一次迁移别一上来就挑战最复杂的模型选中等的练手最稳。2. 迁移前的环境准备与基线验证2.1 可复现的版本组合迁移第一步不是写代码是把版本环境锁死避免导出成功但推理乱码的情况。我这轮用的组合是Python 3.10、PyTorch 2.1.2、transformers 4.36.2、onnx 1.15.0、onnxruntime 1.17.0、protobuf 3.20.3。这里特意点名 protobuf。很多人在导出时遇到constant folding subgraphs之类的诡异报错折腾半天模型代码其实根因就是 onnx 安装时把 protobuf 拉到了低版本或者你自己环境里某个包把它降级了。建议只要做 ONNX 相关开发就把 protobuf 固定到 3.20.3 以上一劳永逸。安装命令很简单pip install torch2.1.2 transformers4.36.2 onnx1.15.0 onnxruntime1.17.0 optimum1.17.1 protobuf3.20.3如果你要跑 GPU 版把 onnxruntime 换成 onnxruntime-gpu版本号保持一致其他不用动。2.2 原模型的加载与冒烟测试导出之前必须先确认一个基线原模型在 PyTorch 下翻译结果正常且你清楚它的分词规则。这一步省了后面所有对比都会变成无源之水。先用最简单的方式加载模型from transformers import MarianMTModel, MarianTokenizer model_name 你下载的模型目录路径 tokenizer MarianTokenizer.from_pretrained(model_name) model MarianMTModel.from_pretrained(model_name) model.eval() text Hello, how are you? enc tokenizer(text, return_tensorspt) result model.generate(**enc) print(tokenizer.decode(result[0], skip_special_tokensTrue))这里有一个很容易被忽略的点Marian 家族不少模型要求源文本前面加语言引导符比如zh这样的格式具体写什么要以模型卡为准。冒烟测试时务必把这个规则确认好因为后面 ONNX 推理循环里要复用同一套前处理否则生成结果会和原模型不一致排查起来极其痛苦。2.3 导出链路选型optimum 还是手写脚本现在导出 ONNX 有三条常见路线optimum 的命令行导出、旧版 transformers.onnx、手动 torch.onnx.export。我的建议很明确优先 optimum然后用手写脚本补齐理解。旧版 transformers.onnx 不要用了它已经不在维护列表里而且对较新的 transformers 版本兼容性不好。optimum 算官方推荐的现代方案内部帮你做了很多算子兼容和配置处理导出一次性成功率最高。但它也有问题封装层级太高一旦出问题你很难定位到具体是哪个前向函数被 trace 失败。所以我会在下一节同时讲两条路先用 optimum 快速出产物再用手写脚本理解它是怎么把 encoder 和 decoder 拆成两个图的。真正上了生产你大概率最后还是需要手写脚本因为业务逻辑往往要自定义输出名、固定 batch、或者裁剪冗余输出。3. 核心实操导出 encoder 和 decoder 两个 ONNX 图3.1 用 optimum 一键导出的完整命令与产物说明先看最快路径。假设模型已经下载到本地./model_dir执行optimum-cli export onnx \ --model ./model_dir \ --task sequence2sequence-lm \ --opset 14 \ onnx_output/task 参数这里必须用sequence2sequence-lm不能随手写成text-generation或者feature-extraction否则导出的图不完整。执行成功后onnx_output/下会出现两个核心文件encoder_model.onnx源语言编码器输入input_ids和attention_mask输出last_hidden_state。decoder_model.onnx目标语言解码器输入包含decoder_input_ids、decoder_attention_mask、encoder_hidden_states、encoder_attention_mask输出 logits。这轮产物里没有所谓的“完整模型单一文件”。Encoder-Decoder 架构在 ONNX 下通用的做法是拆成两个图由外部脚本负责它们之间的数据流动。有一点要提醒optimum 导出后的 decoder 图不包含 HuggingFace 原来的generate逻辑也不包含采样策略所以别指望拿 onnxruntime 直接调用一个run就拿到翻译结果。它只是把模型计算固化成图生成循环还得自己写。3.2 手写导出脚本两种方式都跑一遍optimum 跑通之后我强烈建议再手写一遍导出不是为了重复造轮子而是为了掌握两个关键点动态维度怎么命名、解码器怎么单独包装。后面写推理循环时这些名字必须完全对上。先导出 encoder。把 model.get_encoder() 包装一下保证 forward 签名简洁import torch from transformers import MarianMTModel, MarianTokenizer model_name ./model_dir tokenizer MarianTokenizer.from_pretrained(model_name) model MarianMTModel.from_pretrained(model_name).eval() hidden_size model.config.d_model bos_id tokenizer.bos_token_id or model.config.decoder_start_token_id class EncoderWrapper(torch.nn.Module): def __init__(self, encoder): super().__init__() self.encoder encoder def forward(self, input_ids, attention_mask): return self.encoder(input_idsinput_ids, attention_maskattention_mask)[0] enc_wrapper EncoderWrapper(model.get_encoder()) dummy_input_ids torch.tensor([[1, 2, 3]], dtypetorch.long) dummy_attn torch.ones_like(dummy_input_ids) torch.onnx.export( enc_wrapper, (dummy_input_ids, dummy_attn), encoder_model.onnx, input_names[input_ids, attention_mask], output_names[last_hidden_state], dynamic_axes{ input_ids: {0: batch, 1: seq}, attention_mask: {0: batch, 1: seq}, last_hidden_state: {0: batch, 1: seq}, }, opset_version14 )注意dummy_input_ids的长度取 3 就行动态维度展开后任何长度都能跑但如果你某些算子实现得比较死trace 时会把长度常量固化那就要在导出后多测几个长度。decoder 需要多包一层因为输入不光有 decoder 自己的 token 序列还要接收 encoder 的输出class DecoderWrapper(torch.nn.Module): def __init__(self, model): super().__init__() self.decoder model.get_decoder() self.lm_head model.lm_head def forward(self, decoder_input_ids, decoder_attention_mask, encoder_hidden_states, encoder_attention_mask): out self.decoder( input_idsdecoder_input_ids, attention_maskdecoder_attention_mask, encoder_hidden_statesencoder_hidden_states, encoder_attention_maskencoder_attention_mask, use_cacheFalse )[0] return self.lm_head(out) dec_wrapper DecoderWrapper(model) dummy_decoder_input_ids torch.tensor([[bos_id]], dtypetorch.long) dummy_decoder_attn torch.ones_like(dummy_decoder_input_ids) dummy_encoder_hidden torch.zeros((1, 1, hidden_size), dtypetorch.float32) dummy_encoder_attn torch.ones((1, 1), dtypetorch.long) torch.onnx.export( dec_wrapper, (dummy_decoder_input_ids, dummy_decoder_attn, dummy_encoder_hidden, dummy_encoder_attn), decoder_model.onnx, input_names[ decoder_input_ids, decoder_attention_mask, encoder_hidden_states, encoder_attention_mask, ], output_names[logits], dynamic_axes{ decoder_input_ids: {0: batch, 1: dec_seq}, decoder_attention_mask: {0: batch, 1: dec_seq}, encoder_hidden_states: {0: batch, 1: enc_seq}, encoder_attention_mask: {0: batch, 1: enc_seq}, logits: {0: batch, 1: dec_seq}, }, opset_version14 )这轮踩坑的重点是use_cacheFalse。Marian 模型的 decoder 前向默认会生成past_key_values如果导出时没关掉ONNX 图的输入输出会多出一堆 KV cache 相关张量。虽然 KV cache 能显著加速自回归但在第一次迁移时最好先关闭用最朴素的循环跑通再追求性能优化。另一个注意点是encoder_hidden_states的 dummy 值随手填了torch.zeros。因为它是动态输入所以在 trace 时只关心 shape 和 dtype不关心具体值真正的值会在运行时从 encoder session 的输出传进来。但 dtype 必须固定为 float32不能写成 float64否则导出后推理会报 dtype mismatch。3.3 opset 选择与常见导出报错opset 版本我直接建议 14 起步。ONNX Runtime 1.17 对 14 的支持已经非常成熟大多数 Transformer 算子都能覆盖。如果你发现导出时报某个算子不支持不要急着改模型结构先尝试把 opset 升到 16 或者 17。比较典型的报错是RuntimeError: Exporting the operator ... to ONNX opset 11 is not supported。看到这种提示不要慌检查一下你有没有在 torch.onnx.export 里显式写opset_version。有些封装函数会默认 opset 11而 Marian 的结构在旧 opset 下确实容易碰壁。如果报错信息里出现prim::If或者Unsupported node: If这说明 trace 时走进了模型里的某个动态分支。这个分支一般是 padding mask 或者特殊 token 判断带来的处理办法是确保 dummy 输入 shape 合理且 attention_mask 全部为 1让 trace 过程不进入负分支。实在绕不过去就退回 optimum 导出它内部对你的模型做了不少分支规避。3.4 导出后第一步检查用 onnxruntime 跑一次空数据导出完成后不要直接接业务先用 onnxruntime 检查图的输入输出很多问题都在这一看之下暴露import onnxruntime as ort import onnx onnx.checker.check_model(onnx.load(encoder_model.onnx)) onnx.checker.check_model(onnx.load(decoder_model.onnx)) enc_sess ort.InferenceSession(encoder_model.onnx, providers[CPUExecutionProvider]) for inp in enc_sess.get_inputs(): print(inp.name, inp.shape, inp.type) for out in enc_sess.get_outputs(): print(out.name, out.shape, out.type)这里要注意检查和推理用的是同一个库。如果你只装了 onnxruntime-gpu那就必须写providers[CUDAExecutionProvider]名称拼错会直接报错这也算一个常见低级坑。输出里 shape 出现None是正常的那代表动态维度比如 batch 维和 seq 维。但如果固定维出现问题比如 decoder 的encoder_hidden_states第三维不是你模型的 hidden size那多半是 dummy 输入设置错了排查方向就是回看hidden_size取值和 config 是否一致。4. ONNX 推理循环与结果对齐4.1 从零写一个贪心翻译循环导出成功后最花时间的是推理循环。这里先给一个最简单、最稳的贪心解码实现目标是不求快、只求结果正确。import numpy as np import onnxruntime as ort enc_sess ort.InferenceSession(encoder_model.onnx, providers[CPUExecutionProvider]) dec_sess ort.InferenceSession(decoder_model.onnx, providers[CPUExecutionProvider]) def translate(text, max_len64): tokens tokenizer(text, return_tensorsnp) input_ids tokens[input_ids].astype(np.int64) attn tokens[attention_mask].astype(np.int64) enc_out enc_sess.run(None, { input_ids: input_ids, attention_mask: attn })[0] dec_input np.array([[bos_id]], dtypenp.int64) for _ in range(max_len): dec_attn np.ones_like(dec_input) feed { decoder_input_ids: dec_input, decoder_attention_mask: dec_attn, encoder_hidden_states: enc_out, encoder_attention_mask: attn.astype(np.longlong), } logits dec_sess.run([logits], feed)[0] next_id int(np.argmax(logits[0, -1])) if next_id tokenizer.eos_token_id: break dec_input np.concatenate([dec_input, [[next_id]]], axis1) return tokenizer.decode(dec_input[0], skip_special_tokensTrue)这个循环的核心逻辑是先跑一次 encoder 拿源语言向量然后从 BOS 开始每步把当前已生成的所有 token 喂给 decoder取最后一个位置的 logits 做 argmax把新 token 拼进输入序列直到碰到 EOS 或达到 max_len。我在这里想专门提醒一句encoder_attention_mask的 dtype 必须和导出时 dummy 输入一致。如果导出时是 int64运行时给 float32 会导致 decoder 内部算子报错之前就有同事在这上面耗了半天最后发现就是把astype(np.longlong)写成了astype(np.float32)。这个朴素循环的缺点也很明显序列越长decoder 输入越长计算量是二次增长的。因为每一步都把过去所有 token 重新算了一遍完全没有使用 KV cache。但对于首次跑通验证来说这个实现是最安全、最容易对照原模型的。4.2 与原模型输出做一致性校验跑出第一句翻译后先别高兴拿同一句话和 PyTorch 原模型对比一遍。对比时要把原模型的生成参数固定成和 ONNX 循环一致result model.generate( **enc, do_sampleFalse, num_beams1, max_new_tokens64, )重点对比最终翻译文本而不是 logits。ONNX 和 PyTorch 内部有浮点精度差异哪怕同一个算子不同 kernel 实现也可能让 logits 差一个很小的量进而导致 argmax 在某个 token 上不同。翻译结果大致一样就不用慌这是正常现象。如果差得离谱按顺序排查三件事第一encoder 输出有没有正确传给 decoder第二decoder 的初始输入是不是 BOS第三attention_mask 的 dtype 和 shape 对不对。八成问题都在这三处。另外建议建一个 20 句左右的小测试集涵盖短句、长句、含数字的句子、含换行的文本把 ONNX 和 PyTorch 的输出都保存下来对比。每次改完导出脚本跑一遍回归能省下很多心理阴影。4.3 性能实测与优化手段朴素循环跑通之后接下来就是性能调优。我这轮实测下来几个手段按收益排序如下。第一是固定序列长度。把 decoder 的dec_seq动态维改成静态值比如直接设为 64能省掉 shape inference 的开销。具体做法是导出时把 dynamic_axes 里decoder_input_ids的 seq 维去掉并确保每次推理都 padding 到同一长度。这个改动在短语句场景下体感提升接近两成代价是长句会被截断需要业务上接受 max_len 限制。第二是设置 onnxruntime 的线程参数。默认情况下 onnxruntime 会使用 CPU 的所有核心但线程数不是越多越好线程切换开销会吃掉性能。我一般用session_options.intra_op_num_threads设置为 vCPU 数量的一半左右再配合session_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL通过减少线程竞争获得稳定性能。第三是预热。服务启动后第一次推理会包含模型加载、线程池创建、内存分配等一次性开销首次延迟能比平时高好几倍。上线前在后台异步跑一次空翻译让 session 完成预热是性价比最高的优化。量化这里多说一句很多人在 CPU 部署时急着上 INT8 动态量化但对翻译模型来说中间层误差会随着自回归步数累积很容易出现整段乱码。我实测下来INT8 动态量化在 Marian 上的成功率并不乐观强烈建议先跑 float32实在内存吃紧再考虑精度下降的取舍。5. 常用问题排查清单5.1 常见问题的排查对照表为了让后面接手的人少走弯路我把这轮遇到的高频问题整理成了一个对照表。现象可能原因解法导出报 constant folding 错误onnx 把 protobuf 拉低了固定 protobuf 为 3.20.3 以上报 Exporting operator 不支持opset 版本太低将 opset 升到 14、16 或 17decoder 输入 shape 对不上动态维度命名不一致打印 session 的 get_inputs 逐个比对翻译结果和原模型差异大attention_mask dtype 不对确保运行时 dtype 和导出 dummy 一致第一句话特别慢后续正常缺少预热机制服务启动时跑一次空翻译CPU 并发越高反而越慢线程数开太多设置 intra_op_num_threads 为 vCPU 一半INT8 量化后输出乱码量化误差随自回归累积换回 float32或用固定长度减负5.2 几个值得展开的排障案例这里展开两个我印象最深的案例。第一个是导出后的 decoder 输入顺序问题。optimum 导出的图输入顺序是固定的而手写导出的图输入顺序取决于你在input_names里的排列。如果在推理时想当然地按名字传字典一般没问题但如果你习惯按位置传dec_sess.run(None, [a, b, c, d])那就特别容易踩中顺序错位。排查手段就是打印dec_sess.get_inputs()得到官方顺序再和你的 feed 字典逐项对照。这是所有 ONNX 迁移里最花时间的隐形坑。第二个是 EOS 永远不会触发。我一开始跑的时候循环经常超过 max_len 被硬截断翻译结果末尾多一串没有意义的 token。后来发现是 Marian 某些 checkpoint 的eos_token_id和模型在 generate 时实际使用的 EOS 不是同一个导致我的循环拿着 tokenizer 默认值去判断永远匹配不上。解决办法是读取generation_config或者模型 config 里的eos_token_id而不是默认用一个 0 去猜。这个细节在 HuggingFace 版本升级后尤其容易中招。还有一个小坑也值得一提decoder_input_ids的 dtype。ONNX 对输入类型要求很严格如果你不小心把它喂成 float32运行时不会立刻报错而是产生一个看似正常的乱码结果这个比报错还难查。我在循环开始时显式astype(np.int64)就是为了杜绝这类问题。6. 最后一点实战体会回到开头那个朋友的问题我最后给他的建议很直接如果生产环境允许先别急着把最大模型搬上去。选一个中等尺寸的 MarianMTONNX 导出后用贪心解码加固定长度单机 CPU 也能扛住中等并发。迁移这件事最大成本从来不是导出成功而是导出之后谁来维护推理循环。我自己的体会是写一个小的模型服务层把 tokenizer 规则、session 创建、翻译循环都封装成配置化模块后续换模型只改配置不动代码会省掉后面无数坑。再分享一个小技巧不要把 tokenizer 试图打进 ONNX保留原来的分词器文件在服务端用同样的方式初始化即可因为 tokenizer 的分词逻辑远比前向计算复杂强行静态化只会自找麻烦。先跑通再优化这是这次迁移项目最真实的经验。
