前阵子接了个线上部署的活儿要把一个英译中的功能从实验脚本变成稳定的服务。原来直接用 HuggingFace Transformers 加载模型跑推理单机测试时感觉还挺顺真正并发压上来之后问题就来了模型加载要吞掉不少显存、推理延迟波动大、Python 侧的环境依赖也不太好收敛。排查了半天最后我决定把 HuggingFace 上的这个英译中模型迁移到 ONNX用 ONNX Runtime 来做推理。现在整套流程已经跑通也沉淀出了一套可以复用的迁移方法。如果你也在做 NLP 服务的落地部署手里刚好有一类基于 Transformer 的翻译模型这篇文章应该能帮你省下不少自己踩坑的时间。我用的模型是很常见的 MarianMT 英译中模型架构上是标准的 encoder-decoderHuggingFace 上有现成的权重和 tokenizer。这类模型从 PyTorch 直接推理切到 ONNX最核心的点不是把torch.onnx.export跑通就完事而是要解决编码器/解码器的拆分、动态轴声明、past_key_values 的处理、推理循环重写、以及量化验证这一整串问题。下面我把整个过程按实操顺序拆开讲每一步都会说明为什么要这么做有哪些值得注意的细节。1. 这个迁移到底要解决什么问题1.1 直接跑 PyTorch 模型时遇到的几道坎最开始我是在一个内部服务里直接调用transformers的pipeline做翻译。对单条测试来说没有任何问题但服务化之后痛点会非常集中地暴露出来。第一是环境依赖太重。transformers本身带了一堆子依赖再加上 PyTorch 全家桶部署镜像的体积膨胀得很厉害。服务端如果同时跑很多个模型每个进程都要加载一个完整的 PyTorch 运行时内存开销相当可观。第二是推理性能不稳定。在 CPU 上跑 Transformer 模型时PyTorch 的默认推理路径有时不会触发一些图融合优化算子粒度比较碎线程利用率也不高。换到 GPU 上虽然能跑但小请求的延迟和 batch 带来的收益往往对不上服务的启动时间也很长因为 CUDA 上下文初始化的开销在那摆着。第三是并行控制不够灵活。PyTorch 自己的多线程模型在并发场景下容易出现线程池抢占表现为延迟抖动。而 ONNX Runtime 的 session 隔离性做得更好部署时还能指定并发线程数性能特征要稳定得多。所以这个迁移的目标非常明确把整个翻译链路中最重的模型计算部分从 PyTorch 脱离出来变成一个轻量的、可独立部署的 ONNX Runtime 推理模块。1.2 为什么偏偏选 ONNX 而不是其他方案有人会问为什么不直接用 TorchScript或者更进一步直接上 TensorRT、OpenVINO 这种特定硬件优化方案我的实际体会是ONNX 是一个中立的中间表示生态支持面最广。TorchScript 虽然导出简单但它和 PyTorch 的绑定太深跨版本兼容性差换一个部署环境可能就要重新导出。TensorRT 确实能带来很好的 GPU 加速但工程改动量大而且绑定特定 GPU 架构和 CUDA 版本通用性受限。OpenVINO 对 Intel 平台很友好但跨硬件部署时又不如 ONNX Runtime 灵活。ONNX Runtime 本身的优势在于它把图优化、算子融合、量化、线程调度这些事都打包好了Python 接口写起来也简单底层是 C 实现启动和推理性能都很稳定。它既支持 CPU 也支持 GPU还可以配合动态轴处理变长文本这些都是翻译服务落地非常需要的特性。我的选型结论是如果团队没有强力的硬件专项优化需求优先做 ONNX Runtime 的迁移性价比最高坑也相对可控。1.3 迁移范围模型选型和预期收益这次迁移针对的是典型的编码器-解码器翻译模型输入是一段英文文本输出是中文译文。模型结构上包括共享词表 tokenizerBPE 级别子词切分6 层 transformer encoder6 层 transformer decoder解码端使用自回归方式逐步生成 token迁移完成后的预期收益主要有三个模型推理从 Python 动态图切换为静态图执行延迟波动降低模型文件体积变小在没有 PyTorch 的环境里也能跑配合量化可以把 CPU 上的吞吐往下压一截。需要注意的是这里说的“自己的英译中模型”不一定特指某个固定模型而是泛指在 HuggingFace 上下载的、你准备用于自己服务的英译中权重。只要底层是 MarianMT 或类似的 encoder-decoder 结构这套迁移方法就适用。2. 动手前的准备工作摸清模型的底细2.1 环境依赖与版本选择不管什么模型迁移之前先把环境理清楚。我这边用的 Python 版本是 3.9核心依赖版本如下torch2.x 系列transformers4.3x 版本onnx1.14 以上onnxruntime1.16 以上numpy1.24 以上版本这块必须提醒一句transformers的模型 forward 签名在不同版本之间有改动尤其是处理 past_key_values 的细节直接影响导出的代码写法。如果你用的是较新的 transformers 版本MarianMTModel的 decoder 部分可能已经改成了和标准BartModel很接近的接口导出思路是一样的只是参数名要对上。原始的onnx包负责导出和模型检查onnxruntime负责推理。如果之后要做量化还得用到onnxruntime.quantization模块它在onnxruntime主包里有不需要额外安装别的依赖。建议开始导出之前先固定一份 requirements 清单把transformers和torch的版本冻结住。否则一旦升级版本导出的 ONNX 模型可能在结构层面发生细微变化排查起来非常费劲。2.2 确认模型的编码器-解码器结构拿到权重之后不要急着写导出脚本先把模型结构看一遍。加载方式很简单from transformers import MarianMTModel, MarianTokenizer model_id 你的模型ID tokenizer MarianTokenizer.from_pretrained(model_id) model MarianMTModel.from_pretrained(model_id) model.eval()打印model之后重点确认三件事model.get_encoder()是否存在model.get_decoder()是否存在decoder 的 forward 签名里有哪些参数在 MarianMT 这种模型里model本身是MarianMTModel内部包含独立的 encoder 和 decoder。你可以通过model.get_encoder()拿到 encoder 子模块通过model.get_decoder()拿到 decoder 子模块。导出的时候真正要处理的是这两个子模块而不是整个model。这里有一个容易被忽略的点MarianMTModel.forward和MarianEncoder.forward、MarianDecoder.forward的输入输出约定并不完全一致。如果直接对整个MarianMTModel做torch.onnx.export很容易把生成循环里的prepare_inputs_for_generation逻辑带进去导出来的模型图里就会混入很多不该出现的算子常见报错就是aten::multinomial这类。2.3 摸清输入输出和 tokenizer 的细节导出之前把 tokenizer 的行为也确认一下。MarianMT 的 tokenizer 比较特殊它自带source_lang和target_lang的概念但实际上很多模型卡上就是通过一个特殊的起始 token 来区分翻译方向。关键记号是print(tokenizer.pad_token_id) print(tokenizer.eos_token_id) print(tokenizer.bos_token_id) print(tokenizer.get_vocab())通常pad_token_id用于 batch 内 paddingeos_token_id表示句子结束decoder 的起始 token 不一定等于bos_token_id具体要看模型生成时的配置在推理阶段encoder 的输入就是input_ids和attention_maskdecoder 的输入是decoder_input_ids还可能包括encoder_outputs、encoder_attention_mask以及多层的past_key_values。这些 tensor 的名字在导出时要显式声明并在推理代码里保持一致。3. 导出 ONNX把手动挡开成自动挡3.1 为什么要拆成 encoder 和 decoder 两部分导出翻译模型和单塔分类模型不一样它不是一次前向就能出结果的。自回归生成意味着 decoder 会被反复调用每次输入一个 token同时带上之前所有步骤的 key/value 缓存。这种动态循环结构无法在一个静态 ONNX 图里直接表达。一种做法是把整个 encoder decoder 一次性导出生成时每次前向都从 encoder 重新算一遍。这在功能上可行但效率很差尤其翻译长句时计算量会成倍增长。另一种做法是把模型拆开encoder 导出成一个独立的 ONNX 模型输入是一整句源文本输出是编码向量表示decoder 导出成一个独立的 ONNX 模型输入是当前步的 decoder token加上已缓存的 key/value 和编码向量外层用 Python 或 C 写一个循环控制自回归生成过程我这里采用的是第二种方式也是工程上最主流的方案。3.2 导出 encoder 的完整操作先导出 encoder。encoder 的结构简单输入input_ids和attention_mask输出最后一层编码器的隐层状态。代码可以这样写import torch from transformers import MarianMTModel model MarianMTModel.from_pretrained(model_id) model.eval() encoder model.get_encoder() encoder.eval() # 构造一个 dummy 输入shape 不重要因为后续要声明动态轴 dummy_input_ids torch.tensor([[5, 12, 8, 9, 3]], dtypetorch.long) dummy_attention_mask torch.ones_like(dummy_input_ids) torch.onnx.export( encoder, (dummy_input_ids, dummy_attention_mask), encoder.onnx, input_names[input_ids, attention_mask], output_names[encoder_outputs], dynamic_axes{ input_ids: {0: batch, 1: seq}, attention_mask: {0: batch, 1: seq}, encoder_outputs: {0: batch, 1: seq} }, opset_version14, do_constant_foldingTrue, )几个关键点opset_version我用的 14这个版本覆盖了大多数 transformer 算子兼容性也比较好。如果环境里的 onnxruntime 版本比较老用 11 也可以但不建议低于 10。dynamic_axes用来声明动态轴。翻译服务里句子长度千变万化batch 大小也可能变化所以这个声明是必须的。encoder_outputs的动态轴除了 batch 和 seq最后一个维度也可以声明为动态但通常固定成 hidden size可以不声明。do_constant_folding打开会提前把一些常量计算折叠掉减少推理时的开销。导出完之后用onnx.checker校验一下模型结构import onnx onnx_model onnx.load(encoder.onnx) onnx.checker.check_model(onnx_model)正常的话不会报错。3.3 导出 decoder 并处理 past_key_valuesdecoder 的导出比 encoder 麻烦核心原因是past_key_values。在带缓存的解码场景下past_key_values 是一个元组每个 decoder 层里包含两组 tensor一组是 cross-attention 的 key/value一组是 self-attention 的 key/value。具体层数和每个 key/value 的 shape取决于模型配置和 transformers 内部实现。为了让 ONNX 图能用一组明确的输入输出表示这些内容最直接的办法是把所有 past tensor 展开成独立的输入输出而不是作为一个元组。导出“不带 past 的 decoder 初始步”模型decoder model.get_decoder() # 构造 dummy encoder 输出 dummy_encoder_hidden_states torch.randn(1, 5, 512) dummy_encoder_attention_mask torch.ones(1, 5, dtypetorch.long) dummy_decoder_input_ids torch.tensor([[2]], dtypetorch.long) # decoder起始token torch.onnx.export( decoder, (dummy_decoder_input_ids, dummy_encoder_hidden_states, dummy_encoder_attention_mask), decoder_init.onnx, input_names[input_ids, encoder_hidden_states, encoder_attention_mask], output_names[logits], dynamic_axes{ input_ids: {0: batch, 1: seq}, encoder_hidden_states: {0: batch, 1: seq}, encoder_attention_mask: {0: batch, 1: seq}, logits: {0: batch, 1: seq} }, opset_version14, )注意这一步里我没有显式传past_key_values所以导出的模型只适用于解码的第一步。它接收 encoder 的输出返回针对当前 token 的预测 logits。如果你之后要用 beam searchlogits是计算 beam score 的依据。导出“带 past 的 decoder 续写步”模型时需要先看一下当前 transformers 版本对 past 的表示。比较稳的做法是先用一个实际输入跑一次前向把 decoder 接受到的所有past tensor 的名字和 shape 打出来import inspect print(inspect.signature(decoder.forward))然后构造包含 past key/value 的 dummy 输入。以常见的 MarianDecoder 实现为例每层会有四个 tensorself attention 的 past key、past valuecross attention 的 past key、past value。假设模型有 6 层 decoder那么就有 24 个 past 输入。导出代码思路past_present_names [] for layer_idx in range(num_decoder_layers): past_present_names.append((fpast_self_key_{layer_idx}, fpresent_self_key_{layer_idx})) past_present_names.append((fpast_self_value_{layer_idx}, fpresent_self_value_{layer_idx})) past_present_names.append((fpast_cross_key_{layer_idx}, fpresent_cross_key_{layer_idx})) past_present_names.append((fpast_cross_value_{layer_idx}, fpresent_cross_value_{layer_idx})) input_names [input_ids, encoder_hidden_states, encoder_attention_mask] output_names [logits] for name, _ in past_present_names: input_names.append(name) output_names.append(past_present_names_dict[name])构造 dummy 时每种 past tensor 的形状需要匹配真实 shapeself attention 的 key/value(batch, num_heads, max_seq_len, head_dim)或(batch, num_heads, head_dim, max_seq_len)取决于配置和版本cross attention 的 key/value对应 encoder 输出的 seq 长度如果你不确定跑一次完整前向把outputs.past_key_values里的 shape 打印出来照着写就行。我的主力方案是导出一个decoder_init.onnx负责第一步导出一个decoder_step.onnx负责后续所有步。推理时先用 init 模型拿到第一步的 logits 和 past再循环调用 step 模型。这样整个图的输入输出非常清晰排查问题也容易。3.4 导出后的模型检查导出完除了onnx.checker我还会用 ONNX Runtime 跑一个假数据确认能正常输出。import onnxruntime as ort import numpy as np sess ort.InferenceSession(decoder_init.onnx, providers[CPUExecutionProvider]) for inp in sess.get_inputs(): print(inp.name, inp.shape, inp.type)这一步能帮你确认输入名和维度声明符合预期。如果发现某些输入名和你推理代码里准备的名字不一致回到导出脚本调整input_names。4. 用 ONNX Runtime 实现完整翻译推理4.1 设计推理循环的总体流程ONNX 模型只是静态计算图翻译过程的控制逻辑得自己写。完整流程如下tokenizer 编码源文本得到input_ids和attention_maskencoder 前向一次得到encoder_outputsdecoder 用起始 token 前向一次得到第一步 logits以及第一组 past_key_values从 logits 里取 argmax 或采样得到第一个目标 token如果没到终止条件把当前 token、encoder_outputs、上一轮 past_key_values 喂给 step 模型循环执行收集所有生成 token解码成最终文本4.2 编码阶段一次前向拿到语义向量编码阶段的代码很直接import onnxruntime as ort enc_session ort.InferenceSession(encoder.onnx, providers[CPUExecutionProvider]) def encode(source_text): tokenized tokenizer(source_text, return_tensorsnp, paddingTrue, truncationTrue, max_length512) input_ids tokenized[input_ids].astype(np.int64) attention_mask tokenized[attention_mask].astype(np.int64) encoder_outputs enc_session.run( [encoder_outputs], {input_ids: input_ids, attention_mask: attention_mask} )[0] return encoder_outputs, attention_mask注意tokenizer返回的数组类型ONNX Runtime 要求 numpy 数组的类型严格匹配所以显式astype(np.int64)是个好习惯。paddingTrue代表 batch 内部会 pad 到等长动态轴可以正确处理这种情况。4.3 解码循环从第一个 token 到结束符解码循环是这个迁移过程的核心。初始化模型和 step 模型的调用方式如下dec_init_session ort.InferenceSession(decoder_init.onnx, providers[CPUExecutionProvider]) dec_step_session ort.InferenceSession(decoder_step.onnx, providers[CPUExecutionProvider]) def translate(source_text, max_length128): encoder_hidden_states, encoder_attention_mask encode(source_text) decoder_input_ids np.array([[decoder_start_token_id]], dtypenp.int64) # 第一步 outputs dec_init_session.run( [logits, 所有present输出名...], { input_ids: decoder_input_ids, encoder_hidden_states: encoder_hidden_states, encoder_attention_mask: encoder_attention_mask, } ) logits outputs[0] past_values outputs[1:] all_token_ids [] for step in range(max_length): next_token np.argmax(logits[0, -1, :]).item() all_token_ids.append(next_token) if next_token tokenizer.eos_token_id: break decoder_input_ids np.array([[next_token]], dtypenp.int64) # 把 past_values 作为额外输入传入 step 模型 step_inputs { input_ids: decoder_input_ids, encoder_hidden_states: encoder_hidden_states, encoder_attention_mask: encoder_attention_mask, } for idx, past_name in enumerate(past_input_names): step_inputs[past_name] past_values[idx] outputs dec_step_session.run(output_names, step_inputs) logits outputs[0] past_values outputs[1:] return tokenizer.decode(all_token_ids, skip_special_tokensTrue)这里面要注意几个很容易搞错的地方decoder_input_ids必须是(batch, seq_len)的二维数组seq_len一直是 1因为 ONNX 图里需要保持这个维度。logits的第二维是解码序列长度。由于每次只输入一个 token最后的 logits 第二维也是 1直接取[0, -1, :]就能拿到当前步所有词的概率分布。past 输入的顺序必须和导出时声明的input_names一一对应。排序乱掉shape 都不报错但结果会完全错乱。所以导出脚本里的命名顺序一定要记录清楚推理时照着填。4.4 精度与一致性验证跑通流程之后先别急着上线验证 ONNX 推理结果和 PyTorch 原模型是否一致。最简单的验证方法是准备一批英文句子用同一个模型分别在 PyTorch 和 ONNX Runtime 上跑 greedy decoding逐句对比翻译结果。理想情况下结果完全一致或者只有个别 token 不一致。如果出现大面积不一致常见原因有两个解码循环里 token 拼接错误导致上下文信息丢失past_key_values 的处理顺序和 shape 不对我通常还会对比中间层的输出把同一个输入分别用 PyTorch 的 encoder 和 ONNX 的 encoder 跑一次比较encoder_outputs的绝对误差。如果 encoder 一致再逐步检查 decoder 的第一步和后续步缩小问题范围。数值误差在一定范围内是可以接受的但输出 token 级别的差异要控制在极少比例。如果差异超过 10%大概率是逻辑错误不是浮点误差。5. 性能优化和量化尝试5.1 动态轴带来的开销和应对ONNX Runtime 处理动态 shape 时部分图优化没法完全静态化会存在额外的 shape 推导开销。对于翻译服务这种输入长度频繁变化的应用这是绕不开的。一个折中方案是如果线上业务能够把句子长度分段固定可以导出几个不同 seq 长度的静态图比如 64、128、256。推理时根据实际长度选择最接近的静态图这样可以获得一部分性能提升。如果不想维护多个模型文件就保持动态图然后用 ONNX Runtime 的线程设置和 session 选项来压性能sess_options ort.SessionOptions() sess_options.intra_op_num_threads 8 sess_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess ort.InferenceSession(encoder.onnx, sess_options, providers[CPUExecutionProvider])在 CPU 部署时intra_op_num_threads的调优对延迟影响很大。不是线程越多越好要根据 CPU 核数和请求并发量做测试。我自己的环境下8 线程比 4 线程提升有限16 线程反而因为上下文切换开销导致延迟变高。5.2 INT8 动态量化又一次明显提速ONNX Runtime 支持动态量化可以把权重从 FP32 压到 INT8。翻译模型这种计算密集型的模块量化之后模型体积能明显下降CPU 推理延迟也有改善。量化代码from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic(encoder.onnx, encoder_int8.onnx, weight_typeQuantType.QInt8) quantize_dynamic(decoder_init.onnx, decoder_init_int8.onnx, weight_typeQuantType.QInt8) quantize_dynamic(decoder_step.onnx, decoder_step_int8.onnx, weight_typeQuantType.QInt8)量化后要重新验证翻译结果。动态量化对模型质量的影响通常不大但个别敏感权重可能会带来 token 层面的漂移。如果发现某类句子质量退化明显建议只对 encoder 做量化decoder 保持 FP32。5.3 手工对齐 PyTorch 与 ONNX 的输出量化之后再次跑 4.4 节的一致性验证。我一般会看三组数据翻译结果完全相同句子的比例不完全相同但语义接近的比例出现明显错译的比例对于前两组可以接受第三组如果出现尽量回到 FP32 模型。对比时记得把 PyTorch 侧的推理也固定为 greedy decoding关闭采样。transformers的generate默认可能带 beam search 和其他随机性和 ONNX 侧的手写循环不完全等价。为了让对比有意义PyTorch 侧最好也手动写一个 greedy 循环两边逻辑完全对齐再来对比。6. 常见问题速查与排查心得6.1 导出时报 Unsupported ops 怎么处理最常见的是aten::multinomial。出现这个算子说明你导出时把带生成策略的路径引入了计算图。解决方法是不要直接导model.forward改成导model.get_encoder()和model.get_decoder()。严格只走子模块的前向逻辑生成策略留在外层代码里控制。另一种是aten::index_put_或aten::scatter相关报错这类算子常见于带动态更新的操作。MarianMT 在导出时整体上比较顺畅但如果遇到优先检查是不是 transformers 版本太新、加入了更多动态控制流算子。可以把opset_version从 14 降到 11有些算子在旧 op 版本里会有替代路径但反过来也可能因为算子不支持而报错。稳妥起见先用 14不行再试 11。6.2 推理结果乱码或对不上的排查翻译结果完全乱掉的概率事件我遇到过的原因主要是两个。一是past_key_values的拼接顺序错乱。导出时如果没严格按统一顺序整理 past tensor推理时会把 self attention 的 cache 当成 cross attention 的 cache 传进去计算图不报错但语义全错。解决方法是在导出脚本里用列表把past_*和present_*一一对应打印出来推理代码里严格按同样顺序构造输入字典。二是encoder_attention_mask没有传给 decoder。导出时如果漏掉这个输入模型在解码时无法正确屏蔽 padding 位置长句翻译会出问题。检查方法还是回到导出脚本确认input_names里包含encoder_attention_mask。6.3 几个容易忽略但影响很大的小坑第一个是decoder_start_token_id。MarianMT 这类模型里这个 token 通常就是tokenizer.eos_token_id或者某个特殊 token。千万不要默认它是bos_token_id否则第一轮解码就拿不到正确的 token 分布。加载模型后先用model.config.decoder_start_token_id确认一下。第二个是tokenizer输出的attention_mask默认值。在return_tensorsnp下如果不显式指定paddingTrue不同长度的句子没法组成 batch单独推理没问题但一旦要并发拼 batch 就会很痛苦。第三个是 ONNX Runtime 的内存模式。长文本翻译时如果max_length设置得太大past_key_values会不断膨胀。服务端要对输入长度做限制否则线程安全和内存占用都会成为隐患。第四个是模型文件管理。导出之后encoder、decoder_init、decoder_step 三个文件要一起保存。它们之间有隐式约定encoder 输出的维度要能被两个 decoder 模型正确处理。换一个模型权重必须整套重新导出和验证不能混用。这几轮走下来我最深的体会是ONNX 迁移真正的难点往往不在torch.onnx.export本身而在于你对模型结构的理解程度。把 encoder 和 decoder 拆清楚把 past_key_values 的每个维度、每个命名都核对明白后面的推理循环就是水到渠成的事。实际部署之后那个 ONNX Runtime 的翻译服务在 CPU 上跑得比原来的 PyTorch 版本稳不少进程内存也低了一大截。后面如果再换其他结构的翻译模型我大概也会沿用这套拆分导出的思路先把结构摸透再动手迁。