前阵子接了个机器翻译小服务的改造需求很明确模型不能继续跑在 PyTorch 里因为线上环境只有 2 核 4G 内存打包镜像一旦带上 torch 直接多出 2 个 GB冷启动还要吃 600MB 以上内存。翻了一圈最后决定把 HuggingFace 上的英译中模型迁移到 ONNX用 ONNX Runtime 做推理。这条路我前前后后踩了不少坑从模型下载、PyTorch 转 ONNX、再到精度验证和 INT8 量化基本把完整链路走了一遍。这篇就当交底的记录想直接抄作业的照着操作即可想深挖原理的我也把关键节点掰开讲了。1. 为什么要折腾PyTorch 推理在真实部署中到底痛在哪1.1 先算一笔资源账2 核 4G 容器装不动 PyTorch很多人平时用 HuggingFace 的 transformers 写 demo一个pipeline(translation)就能跑起来英译中确实舒服。但一旦放到正经部署环境麻烦就来了。首先是依赖体积。一个最小的 PyTorch CPU 版本安装完之后也要 200MB 左右而且它还会把 numpy、tokenizers、regex、requests 等一堆依赖带进来。如果再配 CUDA 版本那就奔着几个 GB 去了。容器镜像打包时如果基础镜像不是足够精简拉取和发布都很痛苦。其次是内存占用。PyTorch 模型加载后光模型参数就是几百 MB再加上运行时缓存、临时激活值、Python 解释器本身的开销一个翻译模型吃 1GB 内存非常正常。我线上那台 2 核 4G 的机器除了跑翻译服务还有别的模块PyTorch 一加载其他服务直接进入 swap 地狱。第三是推理效率。PyTorch 的 eager 模式在算子调度上有很多 Python 层开销对小批量、短序列的翻译任务来说这种开销占的比重会相当明显。ONNX Runtime 不一样它会先把计算图整体拉平做节点融合和内存规划在很多情况下推理速度比直接跑 PyTorch 快不少。所以把 HuggingFace 模型迁到 ONNX本质上解决的不是“能不能跑”而是“部署得舒不舒服”的问题。这也是我这次愿意专门花时间折腾的核心原因。对比项PyTorch TransformersONNX ONNX Runtime安装体积CPU 版约 200MBonnxruntime 约 5-10MB内存占用模型加载后居高不下动态量化后可明显下降语言绑定基本绑定 PythonC/Python/Java 等都有图优化eager 模式较弱内置算子融合/图优化部署分发需要完整 Python 环境可以只带 .onnx 文件1.2 ONNX Runtime 替我们干了哪些脏活累活ONNX 只是一个中间表示格式真正跑起来靠的是推理引擎。我选 ONNX Runtime 的原因很直接它对 Transformers 系列模型支持得不错HuggingFace 官方生态里也有配套的导出工具省掉很多手写转换的功夫。它核心做了几件事第一是图优化把相邻的算子合并掉。像 Transformer 里常见的 LayerNorm 残差连接如果在 PyTorch 里跑是一串小算子ONNX Runtime 有能力把它们融合成一个算子减少内核启动次数。第二是执行计划缓存同一个模型跑第二次以后很多调度开销会被摊薄。第三是量化支持CPU 上的 INT8 动态量化方案比较成熟这也是后文性能优化部分能落地的关键。而且 ONNX Runtime 本身提供的 InferenceSession 是线程安全的不像 PyTorch 推理时还要特别小心 GIL 和线程问题。做一个常驻的 HTTP 翻译服务时直接用同一个 session 并发推理比自己在 PyTorch 上搞多线程简单太多。2. 模型选型不是越大越好英译中场景我选了这个2.1 首选Helsinki-NLP/opus-mt-en-zh 的取舍HuggingFace 上的英译中模型不少但最经典的“小钢炮”还是Helsinki-NLP/opus-mt-en-zh。这个模型属于 MarianMT 系列是 OPUS 语料上训练出来的纯英译中模型不搞花活专精一个方向。它的规模大概在三百多 MB相对动辄几个 GB 的 m2m100 系列轻量很多。因为是单一语言对词表也相对小ONNX 转换后无论是加载时间还是推理延迟都容易接受。实际效果方面日常短句、产品文案、技术描述都能翻得比较自然不会像一些老模型那样“逐字翻译”味太重。当时我也纠结过是不是该上更大的模型刷效果后来冷静想了想线上服务的翻译质量只需要达到“可读、无硬伤”不是要参加学术竞赛。用这个小模型做 ONNX 迁移部署成本和效果之间能取一个不错的平衡。2.2 备选m2m100 系列、NLLB 什么时候用如果你要覆盖的不只是英译中还要做中英、中日、中韩等双向翻译那 MarianMT 这种单一语言对模型就会很笨重。这时候可以考虑facebook/m2m100_418M或facebook/nllb-200-distilled-600M。m2m100 本身是多对多模型一个模型覆盖很多语言但是如果转 ONNX它的输入输出结构会比 MarianMT 更复杂导出文件也更大。NLLB 在低资源语言上表现更好但同样体积不小。我的观点是如果只是要一个稳定的英译中服务没必要杀鸡用牛刀如果业务是多语种中继那再考虑这些大家伙。选择模型时还有一个隐藏点要看模型的 tokenizer 和生成配置是否清晰。比如decoder_start_token_id、eos_token_id、pad_token_id这些在转换 ONNX 后写生成循环时全都要用到。选一个配置完整的模型能省很多事。3. 环境准备镜像站下载模型 依赖版本锁死3.1 国内拉取 HuggingFace 模型的正确姿势HuggingFace 模型权重动辄几百 MB直接访问huggingface.co很容易超时中断。我的解决办法是走 HuggingFace 镜像站。设置一个环境变量就能把下载地址切过去整个过程不折腾export HF_ENDPOINThttps://hf-mirror.com设置完成后无论是huggingface-cli还是transformers的from_pretrained都会自动从镜像站拉取权重。我自己更习惯用命令先下到本地后面导出时也可以断网使用huggingface-cli download Helsinki-NLP/opus-mt-en-zh --local-dir ./opus-mt-en-zh如果中途下载断了可以给 huggingface_hub 设置下载超时比如export HF_HUB_DOWNLOAD_TIMEOUT60另外建议下载完检查一下目录完整性。常见问题是只下了pytorch_model.bin缺少config.json、tokenizer.json、vocab.json等文件后面加载模型会直接报错。直接看--local-dir下文件是否齐全。3.2 依赖清单一套不那么新但能跑的版本为了防止“昨天还能跑今天更新后崩了”的鬼故事我给这次迁移定了一个比较保守的版本组合torch2.0.1 transformers4.34.1 optimum1.13.2 onnx1.14.1 onnxruntime1.16.3 huggingface_hub0.17.3 sentencepiece0.1.99这里重点说一下sentencepiece。MarianMT 的 tokenizer 依赖这个库很多环境里用它做中文分词时容易漏装结果from_pretrained走到 tokenizer 加载阶段才报错。我习惯一开始就装上省得浪费半小时。版本上不建议盲目装最新版 transformers。我试过从 4.36 升到 4.40optimum导出的模型结构和旧版生成逻辑有细微差异导致 ONNX 推理结果和 PyTorch 结果对不上。所以如果你用的是旧项目尽量锁死版本。4. 快速迁移用 optimum-cli 一步导出 ONNX4.1 为什么建议先用官方工具跑通手动导出 seq2seq 模型不是不行但第一次接触很容易被past_key_values的自回归结构卡住。HuggingFace 官方生态里已经有optimum这个工具它内部封装了模型转换、生成循环适配、动态轴设置等一堆细节。我的建议是先别急着写大段torch.onnx.export先把官方工具跑通拿到一个可用的 ONNX 模型再研究底层原理。optimum-cli是optimum包自带的命令行工具能针对不同类型的模型自动转换。它会把 seq2seq 模型拆成encoder_model.onnx和decoder_model.onnx并处理好生成过程中的缓存结构。4.2 一行命令导出生成的文件是什么直接执行optimum-cli export onnx \ --model Helsinki-NLP/opus-mt-en-zh \ --task translation-en-to-zh \ --opset 14 \ ./marian_onnx命令跑完后./marian_onnx目录下会多出这些关键文件marian_onnx/ ├── config.json ├── decoder_model.onnx ├── encoder_model.onnx ├── generation_config.json ├── sentencepiece.bpe.model ├── special_tokens_map.json ├── tokenizer.json └── tokenizer_config.json最核心的就是encoder_model.onnx和decoder_model.onnx。前者负责把源语言句子编码为上下文向量后者负责一步步生成目标语言 token。tokenizer.json等文件则保证我们能用同一个分词器做 tokenize 和 detokenize。这里有个细节optimum导出的 decoder 模型实际上是一个“单步解码器”。它支持传入已有的 decoder token 序列和缓存状态但不会自己循环生成。真正的循环控制逻辑还是由 Python 端的ORTModelForSeq2SeqLM来完成。4.3 用 ONNX Runtime 加载并验证翻译结果导出只是第一步关键是能用 ONNX Runtime 跑起来。加载方式很简单from optimum.onnxruntime import ORTModelForSeq2SeqLM from transformers import MarianTokenizer, pipeline model_dir ./marian_onnx ort_model ORTModelForSeq2SeqLM.from_pretrained( model_dir, providerCPUExecutionProvider, ) tokenizer MarianTokenizer.from_pretrained(model_dir) pipe pipeline( translation, modelort_model, tokenizertokenizer, src_langen, tgt_langzh, ) print(pipe(Hello world, this is an ONNX translation test.))第一次跑通如果看到中文输出说明整个链路已经没问题了。后面再去抠精度、性能、动态轴这些细节。不过我提醒一句provider参数要看机器环境。如果直接写CUDAExecutionProvider但机器上没有 CUDAfrom_pretrained会在运行时抛错。先在代码里打印ort.get_available_providers()检查一下可用列表再决定用哪个执行提供程序。5. 手动导出的完整教材理解动态轴、Encoder 和 Decoder5.1 搞清楚 MarianMT 的输入输出签名用optimum确实快但只看它导出的黑盒以后遇到定制需求还是不会改。比如你想换 beam search想固定最大生成长度想自己控制 batch这些都需要理解模型内部结构。所以我后来还是手动导了一遍。MarianMT 整体上是标准的 Encoder-Decoder 结构。加载模型后可以看到它有两个子模块model.encoder model.decoder其中encoder的 forward 输入主要是input_ids和attention_mask输出是last_hidden_state形状是[batch, src_seq_len, hidden_size]。decoder的输入则复杂一些需要decoder_input_ids当前已经生成的目标 token 序列、encoder_hidden_states来自 encoder 的输出、encoder_attention_mask源端 padding mask以及可选的past_key_values。输出是一个[batch, tgt_seq_len, vocab_size]的 logits 张量。手动导出时past_key_values是最麻烦的因为它内部包含多层、多组缓存张量而且长度会随生成过程动态增加。如果直接拿整个 decoder 去导出ONNX 图很难做到既灵活又高效。业界常用做法是只导出一个“单步 decoder”也就是每次生成一个 token 时都重新走一遍全序列计算。牺牲一部分速度换来清晰的图和可控的代码。5.2 导出 Encoder静态图里只负责编码Encoder 很直接因为它没有循环依赖。我先构造一个 dummy 输入然后用torch.onnx.export导出import torch from transformers import MarianMTModel, MarianTokenizer # 加载本地模型 model_dir ./opus-mt-en-zh model MarianMTModel.from_pretrained(model_dir).eval() tokenizer MarianTokenizer.from_pretrained(model_dir) # 构造 dummy 输入 sample tokenizer( [This is a test sentence.], return_tensorspt, paddinglongest, truncationTrue, max_length512, ) dummy_encoder_inputs ( sample[input_ids], sample[attention_mask], ) torch.onnx.export( model.encoder, dummy_encoder_inputs, 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, )很多人容易漏掉的点是last_hidden_state的序列维seq也要放到dynamic_axes里。如果不设置ONNX 会认为输出序列长度固定一旦输入 sentence 长度和 dummy 不一致推理时直接报 shape 不匹配。5.3 导出 Decoder难点在于不传 past_key_valuesDecoder 需要包一层 wrapper屏蔽掉past_key_valuesclass DecoderStep(torch.nn.Module): def __init__(self, decoder): super().__init__() self.decoder decoder def forward( self, decoder_input_ids, encoder_hidden_states, encoder_attention_mask, ): out self.decoder( input_idsdecoder_input_ids, encoder_hidden_statesencoder_hidden_states, encoder_attention_maskencoder_attention_mask, use_cacheFalse, return_dictFalse, ) return out[0] # logits注意我特意把encoder_attention_mask也传进来了。这一点非常关键后面精度验证部分会专门说。如果只在 wrapper 里传decoder_input_ids和encoder_hidden_states不传源端 mask那么 encoder 的 padding 位置也会参与 cross-attention导致输出 logits 和 PyTorch 原始模型完全对不上。导出代码# 生成一个初始 decoder token decoder_start_id model.config.decoder_start_token_id dummy_decoder_ids torch.tensor([[decoder_start_id]], dtypetorch.long) # encoder_hidden_states 用一次前向拿到 with torch.no_grad(): encoder_out model.encoder( sample[input_ids], attention_masksample[attention_mask], )[0] decoder_step DecoderStep(model.decoder) torch.onnx.export( decoder_step, (dummy_decoder_ids, encoder_out, sample[attention_mask]), decoder_model.onnx, input_names[ decoder_input_ids, encoder_hidden_states, encoder_attention_mask, ], output_names[logits], dynamic_axes{ decoder_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, )这里decoder_input_ids的序列维是动态的因为每生成一个新 token序列长度就会 1。encoder_hidden_states的序列维同样动态对应不同的源句长度。5.4 在 Python 里写一个 Greedy 生成循环对接 ONNX导出不是终点得能用 ONNX Runtime 跑起来。手动导出后的生成循环不算复杂核心流程是源句 tokenize 后得到input_ids和attention_mask用 encoder session 得到last_hidden_state初始化decoder_input_ids为decoder_start_token_id循环执行 decoder session取最后一时间步的 logitsargmax 得到下一个 token遇到 eos_token 或长度达到上限时停止对decoder_input_ids做 decode得到中文文本。import numpy as np import onnxruntime as ort encoder_session ort.InferenceSession(encoder_model.onnx) decoder_session ort.InferenceSession(decoder_model.onnx) def translate_onnx(text, max_len64): encoded tokenizer( [text], return_tensorspt, paddinglongest, truncationTrue, max_length512, ) input_ids encoded[input_ids].numpy() attention_mask encoded[attention_mask].numpy() enc_out encoder_session.run( None, {input_ids: input_ids, attention_mask: attention_mask}, )[0] decoder_input_ids np.full( (input_ids.shape[0], 1), model.config.decoder_start_token_id, dtypenp.int64, ) for _ in range(max_len): logits decoder_session.run( None, { decoder_input_ids: decoder_input_ids, encoder_hidden_states: enc_out, encoder_attention_mask: attention_mask, }, )[0] next_token_id int(np.argmax(logits[:, -1, :])) if next_token_id model.config.eos_token_id: break decoder_input_ids np.concatenate( [decoder_input_ids, [[next_token_id]]], axis1 ) return tokenizer.decode(decoder_input_ids[0], skip_special_tokensTrue)这个版本逻辑很清晰但每次 decoder 推理都会把整个目标序列重新算一遍没有利用历史缓存的past_key_values所以长句生成会慢一些。如果只是演示原理、搞定小流量场景完全够用。生产环境想提速建议直接用optimum导出的模型或者基于 ONNX 自带的缓存机制再做优化。5.5 手动导出时最容易翻车的几个细节手动导出不是把代码抄一遍就能成功过程中有几个细节必须提。第一model一定要切到 eval 模式否则 dropout 会把图里的随机性也带进去即使是导出阶段也可能出现 warning后续推理结果不稳定。第二dummy input 的长度不要选太长我用的是 16 左右只要能把模型结构走通即可太长的 dummy 会让模型文件里出现很多不必要的常量。第三opset_version不建议用太高的版本14 在 onnxruntime 1.16 上支持得很稳。第四torch.onnx.export后最好立刻用onnx.checker.check_model检查一遍格式问题早发现早处理。import onnx for name in [encoder_model.onnx, decoder_model.onnx]: model_onnx onnx.load(name) onnx.checker.check_model(model_onnx)6. 精度验证不要在翻译结果差不多时就觉得万事大吉6.1 怎么科学对比转换前后的模型输出很多人做完 ONNX 转换拿一两个句子翻一眼看到中文大概通顺就说“没问题”。但翻译是语言生成一两句看不出系统性问题。更稳妥的做法是同时从“文本效果”和“数值层面”做双重校验。文本层面可以准备 50-100 条英文句子分别跑原版 PyTorch 模型和 ONNX 模型然后计算 BLEU 或编辑距离。如果平均编辑距离在几个字以内基本可以接受。不过文本指标受解码策略影响如果 PyTorch 端用的 beam searchONNX 端用的 greedy两者本来就不同。所以更严谨的是先统一用 greedy再比较。数值层面直接比较 logits。同一个源句用同样的初始 decoder token分别跑 PyTorch decoder 和 ONNX decoder统计 logits 的最大绝对误差。这个误差通常应该在1e-3量级。with torch.no_grad(): model.eval() enc_out model.encoder(sample[input_ids], sample[attention_mask])[0] pytorch_logits model.decoder( torch.tensor(decoder_input_ids), encoder_hidden_statesenc_out, encoder_attention_masksample[attention_mask], use_cacheFalse, )[0] onnx_logits decoder_session.run( None, {...} )[0] err np.max(np.abs(pytorch_logits.detach().numpy() - onnx_logits)) print(max abs error:, err)如果误差大于1e-2肯定有问题先查输入张量是否对齐再查架构配置是否一致。6.2 一个隐藏的 pad token 坑忘了传 encoder_attention_mask我在第一次手动导出时wrapper 里只传了decoder_input_ids和encoder_hidden_states没传encoder_attention_mask。结果翻译出来的中文语义完全不对甚至出现汉字乱序。排查过程很痛苦。我一开始以为是动态轴设置错了反复看 distribution、shape 都对。最后把 logits 拉出来逐位置对比发现 decoder 输出的概率分布和 PyTorch 差别极大特别是在源句 padding 的位置上。原因其实不复杂MarianDecoder 内部在做 cross-attention 时需要一个 mask 来屏蔽掉 source padding 位置。PyTorch forward 里默认会从encoder_attention_mask构造这个 mask如果不提供词嵌入里补零的区域也会参与 attention模型相当于看到了一堆“空洞”生成结果自然偏掉。解决方式就是把encoder_attention_mask作为一个显式输入传进 decoder wrapper并在 ONNX 推理时和input_ids、encoder_hidden_states一起传入。传完之后 logits 最大绝对误差立刻从0.8降到了1e-4以下。这个坑在optimum导出的模型里其实已经处理好了但我手动导出时还是踩了一遍。所以如果你的模型转换后效果不对第一优先级检查的就是 attention mask 有没有完整传。7. 性能优化FP16 和 INT8 量化实测对比7.1 FP16GPU 上开CPU 上别瞎开ONNX 模型也可以在导出时把权重转成 FP16但前提是推理引擎和硬件支持。GPU 上开启 FP16 效果立竿见影显存占用降一半推理速度也有提升。导出 FP16 版本时核心是把模型参数转成半精度再把 dummy 输入里的encoder_hidden_states也转成 FP16model_fp16 model.half() with torch.no_grad(): encoder_out_fp16 model_fp16.encoder( sample[input_ids], attention_masksample[attention_mask], )[0] decoder_step_fp16 DecoderStep(model_fp16.decoder)但我不建议在纯 CPU 环境里用 FP16。ONNX Runtime 的 CPUExecutionProvider 对 FP16 算子支持并不完整很多情况下会直接抛错。即使能跑它也未必比 FP32 快反而因为要不断做精度转换引入额外开销。所以我的结论是CPU 部署老老实实用 INT8FP16 是 GPU 用户的快乐。7.2 INT8 动态量化CPU 部署的甜点在 CPU 上ONNX Runtime 最实用的优化就是 INT8 动态量化。它会把模型中的权重矩阵从 FP32 降到 INT8同时保持激活值用浮点计算。相比训练后静态量化动态量化实现简单而且对精度影响小得多适用于翻译这种 token 分布变化较大的任务。对两个 onnx 文件分别量化即可from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( encoder_model.onnx, encoder_model_int8.onnx, weight_typeQuantType.QInt8, ) quantize_dynamic( decoder_model.onnx, decoder_model_int8.onnx, weight_typeQuantType.QInt8, )量化后文件体积能从三百多 MB 降到九十多 MB模型加载速度也快很多。在我那台 2 核 4G 的机器上单条短句的翻译延迟从 FP32 的 150ms 左右降到 100ms 以内。这里有个经验先跑通 FP32再量化成 INT8最后单独验证几组关键句子的输出。不要因为 word error 变了一点就否决量化而是看真实文本可读性。动态量化对 decoder 的 logits 精度影响相对大一点因为 decoder 每一步的误差会被累加。所以量化和验证阶段要格外关注长句和多次生成的稳定性。我在实测中发现INT8 模型在 10 句以内的短句上效果和 FP32 几乎一致30 词以上长句偶尔会出现词语重复但只要在解码时加一个简单的 repetition penalty就能压制大部分问题。7.3 执行提供程序的选择与踩坑记录ONNX Runtime 里执行提供程序决定了模型跑在哪个硬件后端。最常用的是 CPUExecutionProvider 和 CUDAExecutionProvider。加载模型时传入 provider 列表即可sess ort.InferenceSession( encoder_model_int8.onnx, providers[ CUDAExecutionProvider, CPUExecutionProvider, ], )注意 provider 列表的优先级排在前面的会优先被使用。如果你希望有 GPU 就用 GPU没有 GPU 就退到 CPU就可以把 CUDA 放前面CPU 放后面。但坑也在这里。如果机器上装了 onnxruntime-gpu却没有正确安装 CUDA 工具链CUDAExecutionProvider并不会自动出现反而providers里只有 CPU。这种情况下的排查顺序是先检查onnxruntime.get_available_providers()再查看nvidia-smi是否能正常输出最后确认cuDNN版本和 onnxruntime 要求是否匹配。我自己遇到过一个问题同一个 ONNX 模型在仅 CPU 的容器里延迟是 100ms在带 GPU 的环境里反而变成 180ms。原因就是模型太小GPU 每次去取数据、再拷贝结果回来数据传输开销远大于计算时间。所以不要盲目认为 GPU 一定比 CPU 快小模型 低并发场景CPU 往往更稳。8. 常见问题速查表下面把这次迁移过程中遇到的典型问题整理成一个速查表方便以后排查问题现象可能原因解决方案huggingface 下载超时/中断网络不稳定设置HF_ENDPOINThttps://hf-mirror.com并调整HF_HUB_DOWNLOAD_TIMEOUToptimum-cli 导出报错optimum 或 transformers 版本不兼容锁版本使用文中的依赖组合torch.onnx.export 报 past_key_values 相关错误直接把整个 decoder 导出用 wrapper 屏蔽use_cacheFalse只导出单步解码跑 batch1 时 shape 报错动态轴没有覆盖序列维检查 input、output 的 dynamic_axes 是否完整CPU 上加载 FP16 模型报错CPUExecutionProvider 不支持 FP16 算子改为 FP32 或 INT8 模型ONNX 翻译结果明显错乱未传 encoder_attention_mask生成循环和导出时都要显式传入源端 maskINT8 量化后长句翻译有重复词量化误差在 decoder 自回归中累积解码时加 repetition penalty或退回 FP32文件导出成功但推理速度慢每次 decoder 都从全序列重新计算换成 optimum 带缓存的标准模型或引入 past_key_values最后分享一个我自己的使用习惯生产环境我优先信optimum导出的标准模型因为它把缓存和循环都封装好了稳定可靠手动导出更适合做技术验证和对模型的深度定制。如果你只是做内部工具不需要追求极致速度直接 optmimum INT8 动态量化CPU 上跑小流量完全够。这次迁移下来我对 ONNX 这条路的感受是前期转换确实有学习成本但一旦跑通后面部署的轻快感是 PyTorch 给不了的。希望这篇能把同路人走过的坑提前填平想继续往端侧走的话ONNX 转 kmodel 又是另一套流程之后有机会再单独讲。