ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

英译中模型从HuggingFace迁移到ONNX的完整实践与量化加速

英译中模型从HuggingFace迁移到ONNX的完整实践与量化加速 做机器翻译的老哥们应该都感受过这么一种怪圈模型在 HuggingFace 上效果好、指标漂亮可真到“上线”这一步一堆隐性成本就压过来了——环境依赖太重、推理延迟不稳、不同设备要重新装一轮 PyTorch、显卡上的 FP32 模型还白占显存。尤其是英译中这种“兵家必争”的翻译方向线上每多一毫秒延迟都是用户可感知的。我最近把自己微调好的一套英译中模型从 HuggingFace 迁到了 ONNX用 ONNX Runtime 做推理顺带做了 int8 量化整个部署包轻了几倍CPU 推理也明显提速。这篇就把完整流程和踩过的坑一次讲清楚不管你是拿现成的Helsinki-NLP/opus-mt-en-zh练手还是迁移自己 fine-tuning 过的模型都能照着做。模型的“迁移”听起来唬人其实就是一句话把 PyTorch 的权重和计算图转换成一份自包含的 ONNX 文件之后只用 ONNX Runtime 推理不再依赖 transformers 的前向逻辑。但翻译模型和普通分类模型不一样它是 encoder-decoder 结构输出是一个词一个词生成的所以转换成 ONNX 之后还得自己实现解码循环。这是整个迁移里最容易被低估的部分。下面我把从模型选型、环境准备、转换命令到推理代码、量化和排查的全过程拆开讲。1. 模型迁移的整体思路为什么非ONNX不可1.1 ONNX 到底解决了什么问题先说结论ONNX 不是为了让模型变“聪明”而是为了让模型“好养”。PyTorch 模型部署时目标机器上必须有一套匹配的 PyTorch 环境版本、CUDA、系统库稍微对不上模型就起不来。ONNX 文件是一个中间表示格式谁都不依赖只要对方有 ONNX Runtime 就能推理跨平台、跨语言可以在 C、Java、Go 里调用也能塞进移动端和边缘盒子。对英译中模型来说ONNX 还有一个很实际的好处构建了共享封装。模型导出后tokenizer 配置、生成参数、权重结构都固定在一个文件包里不需要每次启动都去 HuggingFace 拉代码逻辑。对于一个要跑在线翻译服务的团队这意味着“模型即产物”CI/CD 里直接拷贝产物就行end-to-end 的部署链路能缩短将近一半。我个人体感最明显的一次是在一台只有 4G 内存的瘦客户端上跑翻译服务PyTorch 全家桶装完内存基本告急换 ONNX Runtime CPU 版本只有几十 MB内存占用直接降了两个数量级。这就是为什么我强烈建议做部署的团队认真考虑 ONNX 路线。1.2 英译中模型的选型与“自己的模型”怎么处理英译中方向常见的 HuggingFace 模型有三类Helsinki-NLP/opus-mt-en-zhMarian 架构体积小、速度快适合中短文本翻译是入门迁移的最好样本。facebook/nllb-200-distilled-600M多语言 NLLB 架构质量高但参数量大迁移后模型文件也更大部署难度明显上升。自己 fine-tuning 的模型比如基于 Marian、M2M100 或 T5 微调过的版本只要架构在 Optimum 的支持列表里迁移路径和现成模型完全一样。有人会问“我的模型存在自己的 HuggingFace 仓库里能直接迁吗”可以。导出命令里的--model参数既支持 Hub 上的模型 ID也支持本地路径。如果你的模型已经保存成 transformers 的标准目录结构直接指向本地路径就能转换。唯一要确认的是 model config 里的architectures字段比如MarianMTModel、M2M100ForConditionalGenerationOptimum 会按这个字段选择导出的图结构。选择模型时需要明确需求如果你是做产品 demo 或者内部工具opus-mt-en-zh就够了转换时间短、调试成本低我自己这套迁移主要用它做基准。如果要接生产级别的新闻、商务翻译NLLB 这类大模型质量更稳但对应的 ONNX 解码环节也更吃内存和算力。建议先用小模型打通全链路再无缝切换到大模型流程是一样的。2. 环境准备与模型获取2.1 转换依赖清单转换本身不需要多高配的机器CPU 就能完成。我通常在一个独立的 Python 3.10 环境里做避免和线上环境互相污染。核心依赖如下pip install transformers torch onnx onnxruntime optimum这里optimum是 HuggingFace 官方的导出工具链它支持 Transformers 架构到 ONNX 的自动映射内部做了很多细节处理。onnxruntime后面推理要用torch和transformers是转换时前向计算的环境。如果你要用 GPU 导出 FP16可以额外装torch的 CUDA 版本纯 CPU 导出 FP32 和后续量化也完全够。我建议顺手装一个新版onnx因为旧版可能不支持某些新算子尤其当你的模型用了较新的注意力实现。注意 Windows 环境下尽量用 Python 3.9 以上否则 Optimum 的依赖容易出幺蛾子。提示如果你只想转换不想和 PyTorch 版本纠缠可以用pip install optimum[exporters]拉取完整导出依赖它会处理额外的优化器依赖。实测下来比单独装一堆包更省心。2.2 模型从哪里拿英译中模型在 HuggingFace 上公开可用的很多。正常流程是让 Optimum 自动下载它会使用 transformers 的缓存机制模型会存到本地的~/.cache/huggingface目录下次转换直接命中缓存不需要二次下载。如果你的网络环境直接连接 HuggingFace 比较慢或者经常下载到一半断掉有几种常规办法使用公共镜像站点下载模型文件配置环境变量后 Optimum 会自动走镜象。用huggingface_hub的断点续传功能模型文件会自动复用已下载的.incomplete分片。手动从网页下载模型权重再放到正确的缓存路径。这些操作只涉及如何更快拿到公开模型文件不改变迁移本身的技术路线。无论从哪个渠道拿到模型转换结果都是一致可复现的。为了可复现性我在实际转换时固定了模型版本。比如# 直接用模型ID转换前可以pin一个revision optimum-cli export onnx --model Helsinki-NLP/opus-mt-en-zhmain ...生产环境建议不要用main这种浮动版本而是把 commit hash 写进项目文档避免权重被人更新后线上翻译结果悄无声息地变了。就算自己 fine-tuning 的模型也最好存到私有模型库并记录 commit这是工程化部署最容易忽略的一环。3. 核心转换操作从 PyTorch 到 ONNX 的完整记录3.1 用 Optimum 一键导出Optimum 提供了命令行工具最简单的转换命令如下optimum-cli export onnx --model Helsinki-NLP/opus-mt-en-zh --task translation-with-past en_zh_onnx/注意--task translation-with-past这个参数的意思是“导出带 past key values 缓存的翻译模型”。为什么要带 past因为解码器生成每个新 token 时理论上要把前面所有 token 重新算一遍如果没有任何缓存会非常慢。带缓存导出后解码器每一轮只算新增的那一步推理速度能快数倍。不带这个参数导出的版本只适合贪心调试线上跑会让人怀疑人生。转换过程中Optimum 会加载模型、跑一遍 dummy input 以确定动态维度然后分别导出 encoder 和 decoder 两个子图。我的opus-mt-en-zh在普通 8 核 CPU 上大约 3 分钟导出完成。如果是大模型可以先把权重从 Hub 下载好再断网转换减少不可控因素。3.2 导出产物长什么样转换完成后en_zh_onnx/目录下会生成这些文件文件作用encoder_model.onnx编码器图对应 source 句子编码decoder_model.onnx解码器图负责逐个生成目标 tokenconfig.json模型原始配置包含层数、头数等tokenizer.json分词器配置推理时必须配套加载generation_config.json生成参数如长度限制、eos tokenvocab.json/special_tokens_map.json词表与特殊 token 映射它没有把 encoder 和 decoder 合并成一个 ONNX 文件因为翻译任务本质是循环调用合并成一个图反而复杂。实际推理时ONNX Runtime 会加载两个 session一个管编码一个管生成。如果你的模型是 M2M100 或者 NLLB目录结构看起来差不多只是 decoder 的输入输出会有多语言的decoder_input_ids等字段道理一样。3.3 验证导出是否成功导出完成后不要急着拿去跑先用一个最小脚本确认两个模型文件能正常加载import onnxruntime as ort enc_sess ort.InferenceSession(en_zh_onnx/encoder_model.onnx, providers[CPUExecutionProvider]) dec_sess ort.InferenceSession(en_zh_onnx/decoder_model.onnx, providers[CPUExecutionProvider]) print(Encoder inputs:) for inp in enc_sess.get_inputs(): print(inp.name, inp.shape, inp.type) print(\nDecoder inputs:) for inp in dec_sess.get_inputs(): print(inp.name, inp.shape, inp.type)看到输入输出 shape 里有dynamic_axes的符号维度比如?是正常的因为句子长度是动态的。如果 session 加载报错最常见问题是onnxruntime版本太旧升级到最新版基本能解决。我自己第一次导出后犯过一个低级错误直接拿 transformers 的model.generate()风格去调 ONNX结果报了一堆 missing input 错误。原因很简单——ONNX 导出后的模型不会自动帮你做 token 拼接和解码循环这些都需要我们自己写推理逻辑。这也是下一节的重点。4. 编码器-解码器推理ONNX Runtime 下自己写解码4.1 为什么不能直接喂整句话普通分类模型是“一段话进去一个标签出来”但翻译模型是自回归生成先输入句子得到编码向量然后反复把已生成的部分送回 decoder预测下一个 token。PyTorch 版因为 transformers 封装得好你只要调一个generate()就行。到了 ONNX 这里所有逻辑都摊在明面上。标准流程是把源句子分词转成input_ids和attention_mask。用 encoder session 跑一次得到last_hidden_state也就是 encoder 表示。把 start token 喂给 decoder得到第一个词和 past key values。之后每一轮把新生成的 token 和 past key values 一起喂给 decoder直到遇到 eos。真正复杂的是第 4 步的 past key values 管理。ONNX 导出时decoder 的输入会带类似past_key_values.{layer}.decoder.key的字段这些字段保存每一层、每一步的注意力缓存。如果你第一次跑不填它们输出端会返回present.{layer}.decoder.key等字段你要把这些输出原样截下来下一次作为输入塞回去。4.2 贪心解码的参考实现下面这套代码是我在项目里用的简洁版去掉了一些批处理优化保留了最核心的缓存逻辑方便你理解每一步在干什么import numpy as np import onnxruntime as ort from transformers import AutoTokenizer model_dir en_zh_onnx tokenizer AutoTokenizer.from_pretrained(model_dir) enc_sess ort.InferenceSession(f{model_dir}/encoder_model.onnx, providers[CPUExecutionProvider]) dec_sess ort.InferenceSession(f{model_dir}/decoder_model.onnx, providers[CPUExecutionProvider]) # 获取模型层数用于构造 past cache num_layers 6 # 按自己的模型config设置 def translate(text, max_new_tokens128): # 1. 分词 inputs tokenizer(text, return_tensorsnp) input_ids inputs[input_ids].astype(np.int64) attention_mask inputs[attention_mask].astype(np.int64) # 2. 编码一次 enc_outputs enc_sess.run(None, { input_ids: input_ids, attention_mask: attention_mask, }) encoder_hidden_states enc_outputs[0] # shape: [1, src_len, hidden] # 3. 初始化 decoder 输入 decoder_input_ids np.array([[tokenizer.eos_token_id]], dtypenp.int64) past None generated [] # 4. 逐步生成 for _ in range(max_new_tokens): feed { input_ids: decoder_input_ids, encoder_attention_mask: attention_mask, encoder_hidden_states: encoder_hidden_states, } if past is not None: for layer in range(num_layers): feed[fpast_key_values.{layer}.decoder.key] past[layer][0] feed[fpast_key_values.{layer}.decoder.value] past[layer][1] dec_outputs dec_sess.run(None, feed) logits dec_outputs[0] # 取最后一个位置的 logits next_token_id int(np.argmax(logits[0, -1, :])) if next_token_id tokenizer.eos_token_id: break generated.append(next_token_id) decoder_input_ids np.array([[next_token_id]], dtypenp.int64) # 更新 past输出中第1个开始是 key/value按层交替排列 new_past [] for layer in range(num_layers): key dec_outputs[1 2 * layer] value dec_outputs[2 2 * layer] new_past.append((key, value)) past new_past return tokenizer.decode(generated, skip_special_tokensTrue) print(translate(Hello, this is a practical guide.))代码里每一步都值得解释。第一轮 decoder 只喂了eos_token_id这是因为 Marian 这类模型将eos_token当作 decoder 的起始符。如果你用的是 NLLB 或 MBart起始符可能是bos_token_id务必看generation_config.json里的decoder_start_token_id。past key values 的索引不是随便猜的我建议在拿到 decoder session 后打印全部输出的 name 和 shape看清顺序再写死循环。不同 Optimum 版本、不同模型架构的输出顺序可能有差异这也是项目里最容易翻车的地方。4.3 Beam Search 这块怎么处理上面代码是贪心解码对翻译质量要求不高时够用但想要 Beam Search 提升效果自己手写会很麻烦因为要同时维护多个候选序列的 past key values并做序列裁剪和归一化。我的建议是不要重复造轮子直接用 ONNX Runtime GenAI 库。ORT GenAI 是微软专门为生成式模型设计的推理库它和 Optimum 导出的产物配合得比较好加载模型后内部会帮你管理解码循环和 beam 搜索。流程大致是import onnxruntime_genai as og model og.Model(en_zh_onnx) tokenizer og.Tokenizer(model) params og.GeneratorParams(model) params.set_search_options(max_length128, num_beams4, do_sampleFalse) generator og.Generator(model, params) generator.append_tokens(tokenizer.encode(Hello, this is a test.)) while not generator.is_done(): generator.compute_logits() generator.generate_next_token() output tokenizer.decode(generator.get_sequence(0)[0])这种方式的好处是省掉手写 past 缓存管理beam search 也能直接用。缺点是 ORT GenAI 对模型文件格式有一定要求。如果公司内部有专门做推理引擎的人直接把 ONNX 模型交给他们上 C 端效果更好Python 版本适合快速验证和交付 demo。5. int8 量化轻量化部署的关键一步5.1 量化能带来什么收益ONNX 模型转出来默认是 FP32体积大、计算开销高。英译中这种对话式场景往往对显存和内存敏感尤其你要把一个 600M 的 NLLB 模型怼到几台小机器上FP32 很不划算。int8 量化在 ONNX Runtime 里的收益我用实际数字给个概念opus-mt-en-zh转换后的 encoder 模型FP32 大约 16MB量化成 int8 后降到 4MB 左右decoder 也从几十 MB 降到十几 MB。推理延迟在 CPU 上通常能快 30% 到 60%瓶颈仍在 decoder 的自回归循环。模型体积变小不只是省磁盘更重要的是缓存加载更快、内存占用更低、容器镜像能做得更小。5.2 动态量化实操ONNX Runtime 提供两种量化动态量化和静态量化。对翻译模型这种输入序列长度不确定的场景动态量化最容易上手不需要准备校准数据集它会运行时根据输入动态计算每个 tensor 的 scale。代码很简单from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( en_zh_onnx/encoder_model.onnx, en_zh_onnx/encoder_model_int8.onnx, op_types_to_quantize[MatMul, Attention, Gemm], weight_typeQuantType.QInt8, ) quantize_dynamic( en_zh_onnx/decoder_model.onnx, en_zh_onnx/decoder_model_int8.onnx, op_types_to_quantize[MatMul, Attention, Gemm], weight_typeQuantType.QInt8, )这里我只量化了MatMul、Attention、Gemm这三类权重密集型算子没有量化LayerNormalization这类对数值敏感的小算子避免翻译质量掉太多。由于 encoder 和 decoder 是单独文件量化也是分开做的。量化完成后把推理代码里的 session 路径换成_int8.onnx即可输入输出接口和原来完全一致。5.3 量化后的质量怎么验证量化省了资源代价是精度损失。英译中的可接受标准没有统一答案但你不能只看 BLEU 值因为翻译用户更关注语句是否通顺、专有名词是否准确。我的经验是准备两套测试集一套是标准翻译测试集的 500 句量化前后对比 BLEU 和 chrF 指标。另一套是业务手工挑的 50 句“刁钻句”比如涉及数字、日期、缩写、口语化表达人工看结果。量化后如果发现某些固定句式翻译走样可以考虑只量化 encoderdecoder 保持 FP32这是一个很实用的“折中方案”。decoder 是整个生成过程的性能瓶颈但它的精度敏感度也更高很多翻译链路里 decoder 的细微信号会直接影响整句质量。我自己在几个项目里都采用“encoder int8 decoder FP32”的组合体积和速度损失不大质量基本不掉。刻意提醒一句int8 量化后的模型要和 tokenizer、generation_config 配套使用不要只拷走.onnx文件否则生成时缺 token 映射结果会非常诡异。最好把整个en_zh_onnx目录作为一个部署单元来发布。6. 常见问题与避坑笔记6.1 “invalid feed dictionary”或者 missing input 报错这是刚上手 ONNX 推理时最常踩的坑。原因多半是输入字段名没对齐。Optimum 导出的模型decoder 输入通常叫input_ids、encoder_attention_mask、encoder_hidden_states与一长串past_key_values.*不会自动帮你补。解决办法是把 session 的输入 name 先打印出来按实际的名字组织 feed dict不要凭感觉猜。6.2 生成的句子全是同一个词或者提前结束常见原因有三decoder 起始 token 不对Marian 用 eos 起头M2M100 可能用lang_code或 bos看generation_config.json。past key values 的传递顺序错了把某一层的 key 塞给了另一层。生成的 token id 没有逐步拼进 decoder 输入导致模型永远只看到单 token。排查时先在生成循环里打印每一步的 token id 和对应字符确认序列是在增长还是死循环。6.3 量化后部分长句翻译质量崩掉量化对长句更敏感因为长句的注意力分数分布更广int8 的精度可能装不下细小的差异。如果长句质量下降明显建议把量化范围缩小只量化 encoder或者改用 per-channel 量化方式。在quantize_dynamic里可以通过per_channelTrue开启逐通道量化在很多场景下可以挽回一些精度。6.4 ONNX 推理速度反而比 PyTorch 慢这种情况也有一般是两个原因一是没有启用插值优化和图优化二是解码实现里没有缓存 past key values。还有一个很容易忽略的是onnxruntime的线程数设置。在服务端场景默认线程数可能和容器 CPU 配额不匹配手动设置会话选项能显著改善options ort.SessionOptions() options.intra_op_num_threads 4 sess ort.InferenceSession(decoder_model_int8.onnx, sess_optionsoptions, providers[CPUExecutionProvider])别盲目调高线程数小模型开太多线程线程切换的开销比计算本身还大。我一般先压测不同线程数再确定线上的值。6.5 模型文件损坏或下载不完整HuggingFace 下载大文件时断点续传出问题是常事典型表现是转换时卡在Downloading ...或者加载权重时报file not found。删除缓存目录里的.incomplete文件重新下载通常就能解决。如果是公司内网环境建议提前拉好模型后离线转换避免每次构建都被网络波动干扰。最后的实战心得整个迁移流程走下来我心里最深的体会是ONNX 转换只是起点真正考验人的是“推理循环的设计”。翻译模型从不该被当成一个黑盒单图模型它天生就是“编码器 循环生成器”的组合。理解了这一点后面无论是加 beam search、接 ORT GenAI还是做流式翻译都能顺着同样的思路扩展。我在实际使用中的一个小技巧是把推理代码封装成两个函数——encode_sentence()和decode_step()前者只跑一次后者是循环调用。这样调试时只需盯住 decode 那部分量化切换也方便改会话路径就能从 FP32 换成 int8。如果你要部署到移动端或者嵌入式设备下一步可以把 ONNX 再转成边缘平台自己的格式大部分边缘推理框架都支持直接读取 ONNX 作为中间输入迁移链路是通的。如果只是想在服务里拿到一个可靠又好维护的英译中能力上面的方案已经够用。别急着上大模型先用小模型跑通整条流水线再按业务需要升级模型这会省掉很多不必要的精神内耗。
返回列表