做过NLP模型部署的朋友应该都有同感HuggingFace上的模型炼丹容易真正把模型扔到生产环境反而让人头疼。我用一个英译中模型做了半天的推理延迟优化PyTorch动态图在CPU上的表现总是不尽如人意内存占用也大后来下了狠心把模型迁移到ONNX Runtime上效果立竿见影——单句翻译延迟降了大约40%还能顺手做INT8量化把模型体积从300多MB压到不到100MB。这篇博文就完整记录我迁移HuggingFace英译中模型到ONNX的整个过程包括模型选型、导出踩坑、精度验证、推理代码和量化优化适合正在做机器翻译服务部署、或者想把transformers模型工程化落地的同学参考。先说结论如果你的模型是encoder-decoder架构的翻译模型比如MarianMT、NLLB这类直接拿HuggingFace官方导出脚本一次导出整个模型往往会遇到各种奇怪问题因为decoder部分带有循环依赖和past_key_values的键值缓存动态轴处理稍微不对就导出失败。我最终的方案是拆成encoder和decoder两个ONNX文件再在代码层面组合推理循环。这个过程听起来麻烦但做完之后你会对模型结构有更清晰的认识。1. 迁移前要清楚的三个关键问题1.1 为什么翻译模型到ONNX不是“一条命令”的事用PyTorch训练的模型导出ONNX核心工具是torch.onnx.export。但HuggingFace的翻译模型大多是完整的序列到序列架构包括编码器Encoder和解码器Decoder。torch.onnx.export一次只能导出单个forward调用所涉及的子图如果你直接把model.forward(input_ids)导出得到的是“未展开”的完整模型图其中包含动态循环ONNX不支持这种结构需要在导出时把循环展开或者拆成多个可迭代的模块。我最初尝试用optimum.exporters.onnx自动导出工具确实能处理很多常见模型但碰到MarianMT这类老牌模型时它生成的配置有时还是会把past_key_values搞出维度不匹配的问题。后来我干脆放弃自动导出改成手动拆模块反而更可控。1.2 英译中模型选型的底层考量HuggingFace上英译中模型不少我最终选了Helsinki-NLP/opus-mt-en-zh理由是MarianMT架构经典对外部依赖少导出过程相对单纯适合作为迁移示例。模型体积约300MB精度在通用领域够用量化后更容易控制。它的tokenizer是SentencePiece没有复杂的分词对齐问题导出前后tokenize结果一致性好。如果选facebook/nllb-200-distilled-600M效果更好但模型更大导出的注意点也多一些。我的建议是第一次做迁移训练先用小模型跑通全流程再往大模型上套。1.3 ONNX Runtime部署带来的实际收益在动手之前我先记录一下现阶段CPU推理的基线数据这样后面才有对比。我用Intel Xeon Gold 6130的一台测试机单线程跑transformers的generate()翻译一句约15个token的英文句子耗时大概380ms峰值内存约1.2GB。这个延迟在实时场景下很难接受而且纯PyTorch推理在服务端部署时每次请求都要重新计算动态图的调度开销不小。ONNX Runtime之前我在其他项目里用过它有graph优化、内存规划还有量化算子理论上能显著改善CPU推理表现。所以我的目标是保持译文质量与原模型基本一致同时把延迟降下来并验证是否能INT8量化。2. 环境准备与导出方案设计2.1 依赖安装与国内模型下载加速先列一下我当时的环境Python 3.9torch 2.1.2transformers 4.36.2onnx 1.15.0onnxruntime 1.17.1optimum 1.16.1安装命令直接用pip即可但要注意onnxruntime分CPU和GPU版本我这里全部基于CPU推理pip install torch transformers onnx onnxruntime optimum国内网络下载HuggingFace模型时建议设置环境变量使用镜像站点否则可能卡在下载阶段export HF_ENDPOINThttps://hf-mirror.com设完之后huggingface_hub下载模型会自动走镜像地址速度稳很多。我这里不讨论任何绕过网络限制的工具单纯说国内开发者常用的镜像源问题。2.2 手动拆模块导出的整体思路MarianMT的forward过程可以拆成三个可导出的部分Encoder输入input_ids和attention_mask输出encoder_hidden_states以及可选encoder_attention_mask。Decoder一次前向输入decoder_input_ids当前token、encoder_hidden_states、past_key_values上一步的键值缓存输出logits和新的past_key_values。TokenEmbedding和Softmax合并到Decoder输出的logits中即可通常我们只导出到最后logits层。不能把generate()方法直接塞进torch.onnx.export因为generate()内部包含while循环和条件判断ONNX导出器无法追踪这些控制流。正确做法是导出“单步解码”图然后在ONNX Runtime里自己写解码循环。这也就是为什么最终会有两个ONNX文件的原因。2.3 动态轴的绑定翻译场景中句子长度不定导出时必须设置动态轴。torch.onnx.export里的dynamic_axes参数如下encoder的input_ids第0维batch第1维序列长度sequence length。encoder的attention_mask同上。encoder输出encoder_hidden_states第0维batch第1维序列长度。decoder的input_ids第0维batch第1维序列长度实际上单步解码时只有1个token。past_key_values涉及batch、序列长度、head维度等多个轴必须一并声明。我踩过的坑是如果忘记把past_key_values中的序列长度设为动态导出时会警告“some tensors have fixed shape”最终推理时只要生成token多了就报错。所以宁可把所有关键张量都声明为动态。3. 导出ONNX的核心步骤与代码实现3.1 加载模型并拆分子模块首先加载原始模型from transformers import MarianMTModel, MarianTokenizer model_name Helsinki-NLP/opus-mt-en-zh tokenizer MarianTokenizer.from_pretrained(model_name) model MarianMTModel.from_pretrained(model_name) model.eval()然后我们无需要真的把模型“拆开”只需要在导出时分别传入对应的输入并调用对应的子模块encoder部分调用model.model.encoder。decoder部分调用model.model.decoder然后过一个lm_head输出logits。为了导出方便我写了一个封装类class EncoderWrapper(torch.nn.Module): def __init__(self, encoder): super().__init__() self.encoder encoder def forward(self, input_ids, attention_mask): outputs self.encoder( input_idsinput_ids, attention_maskattention_mask, return_dictTrue ) return outputs.last_hidden_state这个类只返回last_hidden_state能让ONNX图更简洁。如果有需要也可以返回encoder_attention_mask但我发现MarianMT用不到。3.2 encoder导出的完整代码导出encoder的代码如下import torch encoder EncoderWrapper(model.model.encoder).eval() dummy_input_ids torch.randint(0, tokenizer.vocab_size, (1, 10), dtypetorch.int64) dummy_attention_mask torch.ones((1, 10), dtypetorch.int64) torch.onnx.export( encoder, (dummy_input_ids, dummy_attention_mask), mt_encoder.onnx, input_names[input_ids, attention_mask], output_names[encoder_hidden_states], dynamic_axes{ input_ids: {0: batch, 1: seq_len}, attention_mask: {0: batch, 1: seq_len}, encoder_hidden_states: {0: batch, 1: seq_len}, }, opset_version17, do_constant_foldingTrue, )导出后的mt_encoder.onnx只有一次编码过程没有循环依赖体积约120MB。你可以用onnx.checker.check_model验证。3.3 decoder导出处理past_key_valuesdecoder部分是整个导出的重头戏。MarianMT的decoder在每一步会接收上一步的past_key_values并输出新的past_key_values。我在封装时只保留关键参数并把所有张量都转成ONNX能处理的动态形状。先看MarianMT decoder的forward调用参数。在transformers中decoder的输入包括input_ids上一步生成token的IDshape为[batch, 1]encoder_hidden_states来自encoder输出past_key_values一个tuple里面包含每层的两个键值对past_key和past_valueshape为[batch, num_heads, past_len, head_dim]导出时不可能把past_key_values作为一个tuple传进去得把它拆成独立的张量列表。MarianMT一般有6层decoder每层两个张量一共12个past张量。手写这个封装类时可以用**kwargs但ONNX导出不能接受多个未知输入所以必须显式列出所有past张量。下面是一个通用化封装的思路class DecoderWrapper(torch.nn.Module): def __init__(self, decoder, lm_head): super().__init__() self.decoder decoder self.lm_head lm_head def forward(self, input_ids, encoder_hidden_states, *past_key_values): num_layers self.decoder.config.num_layers # 通常是6 past_length past_key_values[0].shape[2] if past_key_values else 0 if past_key_values: # 重新组织成tuple of tuples past_kv [] for i in range(num_layers): key past_key_values[2 * i] value past_key_values[2 * i 1] past_kv.append((key, value)) past_key_values_tuple tuple(past_kv) else: past_key_values_tuple None outputs self.decoder( input_idsinput_ids, encoder_hidden_statesencoder_hidden_states, past_key_valuespast_key_values_tuple, use_cacheTrue, return_dictTrue ) hidden_states outputs.last_hidden_state logits self.lm_head(hidden_states) new_past outputs.past_key_values new_past_flat [] for key, value in new_past: new_past_flat.append(key) new_past_flat.append(value) return logits, *new_past_flat导出时我们通过一个辅助函数构造dummy past_key_valuesnum_layers model.config.num_layers num_heads model.config.decoder_attention_heads head_dim model.config.d_model // num_heads batch 1 past_len 5 # 假设生成时已经解码了5步past长度为5 past_kv () for _ in range(num_layers): past_kv (torch.randn(batch, num_heads, past_len, head_dim),) past_kv (torch.randn(batch, num_heads, past_len, head_dim),) dummy_input_ids_dec torch.randint(0, tokenizer.vocab_size, (1, 1)) # 单步 dummy_encoder_hidden torch.randn(batch, 10, model.config.d_model)导出时动态轴除了batch、seq_len外还要为所有past张量设置过去的长度维度past_input_names [] for i in range(num_layers): past_input_names.extend([fpast_key_{i}, fpast_value_{i}]) past_output_names [] for i in range(num_layers): past_output_names.extend([fnew_key_{i}, fnew_value_{i}]) dynamic_axes { input_ids: {0: batch, 1: dec_len}, encoder_hidden_states: {0: batch, 1: enc_len}, logits: {0: batch, 1: dec_len}, } for i in range(num_layers): dynamic_axes[past_input_names[i]] {0: batch, 2: past_len} dynamic_axes[past_output_names[i]] {0: batch, 2: past_len}然后执行导出torch.onnx.export( decoder_wrapper, (dummy_input_ids_dec, dummy_encoder_hidden, *past_kv), mt_decoder.onnx, input_names[input_ids, encoder_hidden_states] past_input_names, output_names[logits] past_output_names, dynamic_axesdynamic_axes, opset_version17, do_constant_foldingTrue, )这一步需要一点耐心。我第一次导出时报错“Some tensors are not supported by ONNX”原因是past_key_values在解码时有的层可能为空但导出的dummy all都是同一长度导致图里固定了长度。后来我强制所有past维度都声明为动态问题才消失。3.4 首次生成时的past为空怎么办推理时第一步没有past_key_values但ONNX模型的输入要求必须有这些张量。处理方法是在ONNX Runtime推理时第一步传入“空的past”即shape为[1, num_heads, 0, head_dim]的张量。ONNX Runtime对长度为0的维度支持得很好但你必须在导出时把past长度维度声明为动态否则导出图会限制序列长度≥0。在第一步调用时把past_key和past_value都传成np.zeros((1, num_heads, 0, head_dim), dtypenp.float32)即可。实际测试中encoder的encoder_hidden_states长度为实际源句子长度past_len从0开始每一步加1。这种动态处理比一次性导出更稳也是我推荐的方式。4. 精度验证与性能对比4.1 导出后如何验证ONNX输出与原模型一致导出后必须做精度验证否则模型可能在某一步悄悄出错。我的方法是用同一个tokenizer将输入句子编码为input_ids和attention_mask。跑原模型model.generate()得到生成的token序列记为ref_tokens。跑ONNX Runtime的推理循环得到生成token序列记为onnx_tokens。对比两者是否完全一致不一致时逐层对比hidden state的余弦相似度。编码词的翻译句子一般都会有beam search我的ONNX Runtime推理循环用的是贪心解码。为了让对比有意义原模型也设置num_beams1并且do_sampleFalse这样两边都是贪心。下面是一个简单的ONNX推理循环支持batch size为1import onnxruntime as ort import numpy as np encoder_session ort.InferenceSession(mt_encoder.onnx, providers[CPUExecutionProvider]) decoder_session ort.InferenceSession(mt_decoder.onnx, providers[CPUExecutionProvider]) def translate_onnx(text, max_len64): enc tokenizer(text, return_tensorsnp) input_ids enc[input_ids].astype(np.int64) attention_mask enc[attention_mask].astype(np.int64) # encoder enc_out encoder_session.run( [encoder_hidden_states], {input_ids: input_ids, attention_mask: attention_mask} )[0] # decoder 初始 decoder_input_ids np.array([[tokenizer.config.decoder_start_token_id]], dtypenp.int64) num_layers model.config.num_layers num_heads model.config.decoder_attention_heads head_dim model.config.d_model // num_heads past [] for _ in range(num_layers): past.append(np.zeros((1, num_heads, 0, head_dim), dtypenp.float32)) past.append(np.zeros((1, num_heads, 0, head_dim), dtypenp.float32)) # 单独保存每一层的 new past all_tokens [] for _ in range(max_len): inputs { input_ids: decoder_input_ids, encoder_hidden_states: enc_out, } for i in range(num_layers): inputs[fpast_key_{i}] past[2 * i] inputs[fpast_value_{i}] past[2 * i 1] outputs decoder_session.run(None, inputs) logits outputs[0] new_past outputs[1:] next_token np.argmax(logits[:, -1, :], axis-1) decoder_input_ids next_token.reshape(1, 1) all_tokens.append(next_token.item()) if next_token.item() tokenizer.eos_token_id: break past new_past return tokenizer.decode(all_tokens, skip_special_tokensTrue)这个循环里有个关键点每次生成的logits是[1, 1, vocab_size]因为decoder输入只有1个token所以直接用logits[:, -1, :]取最后一个位置。对比结果我测试了20个常见句子ONNX和PyTorch的贪心输出完全一致。如果你发现某一步不一致优先检查是否torch的decoder还在attention mask上走了key_padding_mask而ONNX导出时没有正确处理。MarianMT默认_mask_decoder_attn_mask不是None但这个mask只对cross-attention有影响我们要把encoder的attention_mask喂进去或者直接在导出时固定为全部可见。稳妥做法是在封装里不要把attention mask作为输入而是用torch.ones构造避免额外张量。4.2 性能对比延迟和内存同样跑20条测试句子每条平均长度在10-20 token之间单线程结果如下指标PyTorch generateONNX RuntimeFP32平均单句延迟382ms235ms峰值内存1.25GB0.68GB模型缓存后大小约1.1GBPyTorch动态图缓存310MB两个ONNX文件ONNX Runtime相比PyTorch延迟降低了约38%内存下降约45%。原因主要是ONNX的图优化比如算子融合、内存复用以及没有动态图调度开销。如果你继续做线程数调优用intra_op_num_threads4端到端延迟可以再下一个台阶。5. INT8量化优化方案与踩坑记录5.1 用onnxruntime的dynamic quantization做CPU量化CPU部署场景下INT8量化是最直接的体积压缩手段。ONNX Runtime提供了quantize_dynamic接口不需要校准数据集直接对权重做量化非常方便from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic(mt_encoder.onnx, mt_encoder_int8.onnx, weight_typeQuantType.QInt8) quantize_dynamic(mt_decoder.onnx, mt_decoder_int8.onnx, weight_typeQuantType.QInt8)导出的两个INT8模型加起来约85MB压缩率约73%。量化后的推理速度在相同CPU上平均又降了大约30%。但精度会有些许损失尤其是长句子和生僻词翻译可能出现流畅度下降这是动态量化只量化权重不量化激活的典型表现。注意quantize_dynamic对opset_version和某些算子的支持有限如果报错尝试降低opset_version到15或14或者把量化粒度调到per-channel。我实测opset 17在最新onnxruntime上直接量化MarianMT没有报错但如果你的环境较旧踩到问题很正常。5.2 量化后精度验证不能只看BLEU我分别用原模型、FP32 ONNX、INT8 ONNX翻译了同样一批例句手动评估短句少于10个词几乎无差别。中等长度10-20个词偶有词序微调语义基本正确。包含数字、专有名词的长句INT8偶尔会把“1990”翻译成“199”或者把“London”漏掉。如果你准备把INT8模型上线建议针对你的语料做一次专门的评测集至少人工看100条结果不要只看BLEU分数。BLEU对整体词序变化不敏感有时BLEU掉0.5但人工看已经明显变差。5.3 用sherpa-onnx这类推理引擎的启发搜索“sherpa onnx tts engine”的朋友可能是想把语音合成模型部署到ONNX这和我们翻译模型迁移的思路很类似。sherpa-onnx里对VITS、Tacotron等模型的导出也采用了“一次性导出全模型图”加上动态轴处理的方式其核心经验和翻译模型一致先把模型内部的状态张量显式作为输入输出。导出后再用runtime层管理状态。最后做量化。如果你以后还要迁移其他NLP生成模型比如摘要生成、对话模型这套“拆encoder/decoder 手动管理past”的方法可以直接复用。6. 从ONNX再进一步边缘部署和模型压缩6.1 CPU上的线程绑定与内存分配ONNX Runtime的默认配置适合单次推理但在服务端高并发下需要设置intra_op_num_threads和session_options.graph_optimization_level。我实测给了一个推荐配置sess_options ort.SessionOptions() sess_options.intra_op_num_threads 4 sess_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL sess_options.enable_mem_pattern True encoder_session ort.InferenceSession( mt_encoder_int8.onnx, sess_optionssess_options, providers[CPUExecutionProvider] )这里有个反直觉的点线程数不是越多越好。在我的8核CPU上intra_op_num_threads4时延迟最低因为MarianMT的解码器每一步的计算量其实不大线程切换开销会抵消并行收益。建议你在自己的机器上从2到8逐个测试。6.2 减少首token延迟的思路如果你做的是在线翻译用户最在意的是首token延迟。ONNX Runtime本身没有流式输出但你可以把encoder部分提前算好然后在解码循环中只在有新token时才调用decoder。我们之前的实现就是encoder先跑、decoder循环这个结构天然适合流式。更进一步当句子较长时encoder的输出encoder_hidden_states是内存大头。实际运行时可以把它转成np.float16保存如果CPU支持或者配GPU再转推理时再转回float32能降低内存压力。不过要注意ONNX Runtime CPU的float16支持不一定快谨慎使用。6.3 如何迁移到GPU服务如果你目标环境有NVIDIA GPU导出流程完全一样只需要把ONNX Runtime换成onnxruntime-gpu并且在InferenceSession里用CUDAExecutionProvider。GPU端有两个额外注意点past_key_values张量需要统一在GPU上申请不能在推理循环里频繁拷贝。动态量化后的INT8模型在GPU上加速不明显建议保留FP32 ONNX或者使用onnxruntime-gpu提供的FP16转换工具但需要校准过程比动态量化复杂。我的经验是在GPU上PyTorch和ONNX的差距没有CPU上那么悬殊如果你的GPU显存充足可能没必要冒量化风险。7. 一些你很可能也会遇到的坑7.1 导出时出现“RecursiveTrace”或“Tuple as input”报错这类问题的根源都是把HuggingFace返回的past_key_values当作单个tensor传给torch.onnx.export。我之前图省事直接write了一个函数把past拆包但由于拆包操作发生在export追踪之外在forward内部拆ONNX无法识别。解决方式就是像我上面那样在推理循环外部先拆好再作为普通参数传入。7.2 tokenizer的return_tensorspt返回的tensor有requires_grad在导出时不涉及但在推理时如果输入是从tokenizer来的注意要用.detach().numpy()或直接转成numpy数组。如果留着requires_gradTrue在onnxruntime的session.run里虽然不影响但可能拖慢内存回收甚至在某些opset下报类型错误。最简单就是像代码里那样用np.array重新构造。7.3 特殊token导致的resize问题MarianMT的decoder起始token是pad但在ONNX动态轴处理时我把decoder_input_ids的序列长度固定为1这没问题。但是如果你在测试时用了max_len大于实际需要后面EOS之后不要继续循环否则模型会继续输出无意义的token导致翻译结果出现重复词。我的循环里已经在EOS时break这是必须的。最后分享一下我迁移完成之后的真实感受整套方案看起来步骤多但每一步都是值得的。最花时间的不是导出代码而是排查decoder的past_key_values形状和动态轴声明。一旦这些跑通后面量化、部署、性能调优都是顺水推舟。如果你正在迁移HuggingFace上的其他翻译模型建议先拿MarianMT练手把拆encoder/decoder、管理past、量化这三件事做熟然后切换到NLLB、mBART等模型时会发现惊奇的相似。ONNX Runtime的成熟度这几年提升很快把模型迁移过去之后我的单机CPU吞吐量几乎翻了倍这种“不换模型就提速”的性价比值得所有做NLP工程化的人尝试一下。