★ FEATURED ARTICLE 1. 先搞清楚要迁的是谁Helsinki-NLP/opus-mt-en-zh 到底是个什么模型标题里写“HuggingFace 自己的英译中模型”严格说起来 HuggingFace 自己不训练翻译模型大家默认最常说的这个“HuggingFace 英译中模型”其实是托管在 HF Hub 上的Helsinki-NLP/opus-mt-en-zh。它来自赫尔辛基大学 OPUS 项目是一个标准的 MarianMT 架构的神经机器翻译模型专门做英文到中文的翻译。我这次迁移的主角就是它。选择它有个很现实的原因在 CPU 服务器上做英译中这个模型是最常被拿来用的。参数量在 7000 万到 8000 万左右fp32 权重文件大约 300MB翻译质量对短句和中等长度文本都还不错尤其适合做通用型翻译接口。相比之下facebook/m2m100_418M 虽然多语种能力强但体积大得多T5 系列不是开箱即用的翻译模型还得自己找适配的 checkpoint。所以想快速上线一个英译中能力opus-mt-en-zh 就是默认答案。我一开始的误区是准备直接在服务器上把整个MarianMTModel跑起来反正 transformers 的generate()方法一行代码就能出结果。压测之后发现单条 CPU 推理延迟到了秒级而且为了跑一个翻译接口要在容器里装 torch、transformers、sentencepieceruntime 体积非常夸张。后来才决定把模型转成 ONNX用 ONNX Runtime 作为推理后端。1.1 模型结构拆解Encoder-Decoder 和自回归生成理解这个迁移过程先要把 MarianMT 的结构拆开。它就是一个典型的 Transformer Encoder-Decoderencoder 负责把源语言句子编码成语义向量decoder 负责逐个 token 地生成中文。和 BERT 这类纯 Encoder 模型最大的区别是翻译模型的推理是自回归的每生成一个 token都要把之前所有 token 作为输入再算一遍。具体到这个模型相关结构大概是词表大小 65000 左右6 层 encoder、6 层 decoderhidden size 512分词器是 SentencePiece模型目录里那个source.spm就是分词模型。这在后面导出 ONNX 时是个关键点你不能把整个generate()过程直接塞进一个 ONNX 图里因为 ONNX 没法直接把带 Python 循环的自回归解码过程整个打包。generate()内部用的是 Python 循环控制每一步都要调用一次模型前向计算。如果真的用静态导出强行把generate完整 trace 进去会得到一堆不合理的图结构要么包含固定的循环次数要么把中间变量固化死。1.2 什么情况下值得转 ONNX不是说所有场景都应该无脑转 ONNX。以我自己的判断标准如果你的推理环境是 GPU 且显存充足transformers 的 batch 推理配合 CUDA 依旧是最快的ONNX 的优势不大如果只是本地跑个脚本、翻译几十句话也没必要转但如果你想在 CPU 容器、边缘盒子、无 GPU 服务器上提供服务或者想让部署产物脱离 PyTorch 依赖ONNX 这条路就非常值得走。另外ONNX 后还有一个隐藏收益可以接着做 int8 量化。这是 PyTorch 原版在 CPU 上不太好落地的优化到 ONNX 这里反而简单了。这也是我把整个迁移拆成“先转 ONNX再量化 int8”两段来做的主要原因。2. 准备工作模型下载、环境依赖和翻译基线2.1 国内拉取模型的一个稳妥方式HuggingFace 的原始下载域名在某些网络环境下会超时、断连但这是网络链路问题和模型本身没关系。我的做法是直接用 HuggingFace 官方的国内镜像站点hf-mirror.com通过环境变量切换下载端点。整个过程不需要任何额外工具只需要设置HF_ENDPOINTexport HF_ENDPOINThttps://hf-mirror.com然后使用 huggingface_hub 的snapshot_download把整个模型目录拉下来显式指定本地目录方便后续离线加载import os os.environ[HF_ENDPOINT] https://hf-mirror.com from huggingface_hub import snapshot_download snapshot_download( Helsinki-NLP/opus-mt-en-zh, local_dir./opus-mt-en-zh )下载完成后模型目录里会有config.json、model.safetensors或pytorch_model.bin、source.spm、tokenizer_config.json等文件。这里有一个我踩过的坑一开始我只下了模型权重文件没下source.spm结果MarianTokenizer.from_pretrained直接报错因为 Marian 系列的分词器必须依赖 SentencePiece 模型文件。所以最稳妥的方式永远是snapshot_download整目录拉取不要自作聪明只挑大文件。2.2 建一个干净的运行环境我习惯用虚拟环境避免污染系统 Python。这里给出一个能够复现的依赖清单python -m venv .venv source .venv/bin/activate pip install -U pip pip install torch --index-url https://download.pytorch.org/whl/cpu pip install transformers onnx onnxruntime sentencepiece huggingface_hub把 torch 装成 CPU 版就够用因为导出 ONNX 不需要 GPU当然如果你的机器已经有完整版 torch 也无所谓。onnxruntime 至少要 1.15 以上老版本对 Transformer 里常见算子的兼容性不够好。2.3 先跑通一个翻译基线在动手导出之前我建议先用原始模型跑几组翻译句子把结果记下来。这些结果就是后面比对 ONNX 有没有转错、量化后质量损失多少的“标准答案”。from transformers import MarianMTModel, MarianTokenizer model_dir ./opus-mt-en-zh tokenizer MarianTokenizer.from_pretrained(model_dir) model MarianMTModel.from_pretrained(model_dir) model.eval() test_sentences [ The quick brown fox jumps over the lazy dog., The weather in Shanghai is very nice today, and the temperature is around twenty degrees., He earned 500 dollars last week by working part-time at a coffee shop., ] for text in test_sentences: batch tokenizer(text, return_tensorspt) translated model.generate(**batch, max_new_tokens64) print(tokenizer.decode(translated[0], skip_special_tokensTrue))我这边第一句输出的是“敏捷的棕色狐狸跳过那只懒狗。”这类经典结果。注意这里一定要设置max_new_tokens不设置的话会走旧版max_length逻辑某些版本下可能提前被截断影响后续对比。3. 导出 Encoder 和 Decoder别把整个 generate 流程塞进 ONNX3.1 为什么必须拆开导出前面已经说了generate()是带循环的自回归解码。ONNX 支持通过Loop节点表示循环但要把 transformers 内部的动态循环、条件退出、beam search 都转换成 ONNX 算子工程复杂度非常高而且一旦转换成功后续也很难做量化。更工程化的思路是把模型拆成两个独立的 ONNX 文件。encoder.onnx输入源句子input_ids和attention_mask输出 encoder 最后一层 hidden statesdecoder.onnx输入当前已生成的目标 token 序列、encoder hidden states、encoder attention mask输出下一个 token 的 logits。自回归循环放在 ONNX 图外面的 Python 代码里每一轮调用一次 decoder session。这样每个图都只负责一个前向计算逻辑清晰、容易调试也能分别做量化优化。很多开源翻译项目也都是这么做的。3.2 Encoder 导出动态维度必须显式声明导出 encoder 的代码如下。关键点是dynamic_axes因为源句子的长度是动态的batch 大小如果以后想并行翻译多条最好也设成动态。import torch encoder_model model.model.encoder encoder_model.eval() text The quick brown fox jumps over the lazy dog. inputs tokenizer(text, return_tensorspt) with torch.no_grad(): torch.onnx.export( encoder_model, (inputs[input_ids], inputs[attention_mask]), encoder.onnx, input_names[input_ids, attention_mask], output_names[encoder_hidden_states], dynamic_axes{ input_ids: {0: batch, 1: src_seq}, attention_mask: {0: batch, 1: src_seq}, encoder_hidden_states: {0: batch, 1: src_seq, 2: hidden}, }, opset_version14, )这里有个细节MarianEncoder 的输入其实是input_ids和attention_mask返回的last_hidden_state形状是(batch, src_seq, hidden)。之前有人踩过坑以为 Marian 是旧的(src_seq, batch, hidden)布局结果 decoder 那边怎么拼都报维度错误。实际以tokenizer返回的 tensor 形状为准batch 在前。3.3 Decoder 导出无 KV Cache 版本的取舍Decoder 这一步容易把人绕晕因为有past_key_values这个高阶玩法。先介绍我这次用的简单但完全可用的方案不导出 KV Cache。没有 KV Cache 意味着每一轮生成时decoder 的输入是“已经生成出来的全部 token 序列”。比如已经生成了 3 个 token第 4 轮就把这 3 个 token 一起喂进去得到整个序列的 logits再取最后一个位置的 argmax。这样每轮计算量会随着长度线性增长但优点是导出极简单、几乎不会失败、验证容易。dummy_decoder_ids torch.tensor([[model.config.decoder_start_token_id]]) # 用随机张量作为 dummy encoder hidden states确保它被当成输入而不是常量 dummy_enc_hidden torch.randn(1, 8, model.config.hidden_size, dtypetorch.float32) dummy_enc_mask torch.ones(1, 8, dtypetorch.int64) with torch.no_grad(): torch.onnx.export( model.model.decoder, (dummy_decoder_ids, dummy_enc_hidden, dummy_enc_mask), decoder.onnx, input_names[ decoder_input_ids, encoder_hidden_states, encoder_attention_mask, ], output_names[logits], dynamic_axes{ decoder_input_ids: {0: batch, 1: tgt_seq}, encoder_hidden_states: {0: batch, 1: src_seq, 2: hidden}, encoder_attention_mask: {0: batch, 1: src_seq}, logits: {0: batch, 1: tgt_seq, 2: vocab}, }, opset_version14, )这里最需要注意的是encoder_hidden_states不能直接用 encoder 导出的真实结果传入因为那样在 trace 时可能被当作常量固化到图里。用随机张量做 dummy 输入就能保证它成为 ONNX 图的动态输入节点。导出完成后运行时传入什么值都行只要维度匹配。有些人会问要不要把 KV Cache 也导出来我的建议是先跑通无缓存版把整条链路验证完再去折腾 KV Cache。带 KV Cache 的 decoder 导出光是 past_key_values 相关的输入输出就要展开 12 个张量6 层乘以 key/value调试成本一下子高很多。对大部分业务场景而言“无缓存 int8 量化”已经足够用了。3.4 导出后的第一次验证导出完成后不要急着写生成循环先独立验证两个 ONNX 图和 PyTorch 原模型输出是否一致。import onnxruntime as ort import numpy as np sess_enc ort.InferenceSession(encoder.onnx, providers[CPUExecutionProvider]) ids inputs[input_ids].numpy() mask inputs[attention_mask].numpy() onnx_enc_out sess_enc.run(None, {input_ids: ids, attention_mask: mask})[0] with torch.no_grad(): torch_enc_out model.model.encoder(input_idsinputs[input_ids], attention_maskinputs[attention_mask]).last_hidden_state.numpy() print(encoder max abs diff:, np.abs(onnx_enc_out - torch_enc_out).max())正常情况这个差异应该在1e-5到1e-4量级。如果差异到了1e-1那多半是 dropout 没有关干净或者模型没有eval()。decoder 那边也可以做类似验证用同一个encoder_hidden_states喂给 PyTorch decoder 和 ONNX decoder比对最后一步 logits。上面对齐了生成循环才有意义。4. 用 ONNX Runtime 写生成循环把拆开的模型拼回来4.1 Session 初始化和 Provider 选择导出只是阶段一真正跑起来还要把拆开的两个模型在推理侧重新拼起来。先初始化两个 sessionsess_enc ort.InferenceSession(encoder.onnx, providers[CPUExecutionProvider]) sess_dec ort.InferenceSession(decoder.onnx, providers[CPUExecutionProvider])如果你的机器有 Nvidia GPU 且 onnxruntime-gpu 也装了可以把providers改成[CUDAExecutionProvider, CPUExecutionProvider]ONNX Runtime 会自动选择可用设备。不过本文主要讨论 CPU 场景int8 量化也是主要为了 CPU 加速。这里有一个非常容易忽略的坑ONNX Runtime 对输入 dtype 极度敏感。input_ids必须是int64encoder_hidden_states必须是float32。如果从 numpy 数组转过来时默认生成了int32或float64session 运行时会直接报类型不匹配。所以喂数据之前最好显式转换 dtype。4.2 贪心搜索的完整实现下面这个translate函数是整篇文章最核心的部分。它的流程是先跑 encoder拿到 encoder hidden states再从 decoder_start_token_id 开始循环调用 decoder每次取最后一个位置的 logits 做 argmax遇到 eos_token_id 或达到最大长度就停。def translate(text, max_new_tokens128): batch tokenizer(text, return_tensorspt, paddingTrue) input_ids batch[input_ids].numpy().astype(np.int64) attention_mask batch[attention_mask].numpy().astype(np.int64) encoder_hidden_states sess_enc.run( None, {input_ids: input_ids, attention_mask: attention_mask}, )[0] decoder_ids np.array( [[model.config.decoder_start_token_id]], dtypenp.int64 ) for _ in range(max_new_tokens): logits sess_dec.run( None, { decoder_input_ids: decoder_ids, encoder_hidden_states: encoder_hidden_states, encoder_attention_mask: attention_mask, }, )[0] next_token int(logits[0, -1, :].argmax()) if next_token model.config.eos_token_id: break decoder_ids np.concatenate( [decoder_ids, np.array([[next_token]], dtypenp.int64)], axis1, ) return tokenizer.decode(decoder_ids[0], skip_special_tokensTrue)实际运行起来decoder_ids每一轮都会变长ONNX Runtime 会根据动态轴的设置重新推断图结构第一次调用时会有点初始化开销后面会稳定。如果发现第一次特别慢可以先在服务启动时跑一个空句子做预热效果很明显。4.3 和 PyTorch 原模型逐句对比写好translate函数后我用之前保存的三条测试句子和 PyTorch 基线做了对比。短句的结果基本完全一致长句偶尔会有选词不同但句子含义没有跑偏。如果对比出现大量不一致优先查两件事encoder 输出是否有差异decoder 的 logits 是否因为在导出时把encoder_hidden_states意外固化而出现偏差。我一度遇到一个问题ONNX 输出的短句完全正常但长句从某个 token 开始疯狂重复。后来发现不是因为模型转坏了而是我偷懒没设置最大长度导致循环在 eos 未被触发时一直生成最终退化成了重复循环。所以max_new_tokens一定要设上并且长句场景建议做一次后处理如果检测到连续重复 token可以提前截断。4.4 进阶方向Beam Search 和 KV Cache贪心搜索实现最简单但翻译质量上通常还是 beam search 更好一点。如果想把贪心换成 beam search思路是在 Python 循环里维护多个候选序列每轮把当前 batch 的所有候选序列一起喂给 decoder再用 score 筛选。这种写法并不复杂但会增加大概几十行代码。KV Cache 的优化则是另一条路线。无缓存版每生成一个新 token都要对前面的所有 token 重新计算 self-attention有 KV Cache 的话只需要计算新 token 的 key/value 和 query同时跨 attention 部分也可以复用 encoder 侧的缓存。这个优化对长句子的加速非常明显代价是导出时要把 key/value 缓存张量作为额外的输入输出。我建议等无缓存链路完全稳定后再根据线上延迟指标决定要不要投入这个改造。5. int8 量化让模型再瘦一圈CPU 上跑得更快5.1 为什么是先做动态量化而不是静态量化ONNX Runtime 的量化大体分两类动态量化和静态量化。动态量化不需要提前准备校准数据集权重被量化为 int8激活仍然以 float 计算量化和反量化发生在算子内部静态量化需要准备一批代表真实输入分布的校准数据提前统计激活的 min/max 范围运行时激活也走整数路径。翻译模型的输入长度动态性很强句子长短差异大如果搞静态量化校准集很难覆盖真实分布。一旦某个激活值超出统计范围翻译质量会明显劣化最典型的表现就是长句输出崩坏。所以我选择了动态量化省事且稳定。对翻译任务来说动态量化在 CPU 上的加速收益已经很可观完全足够用于生产。5.2 量化脚本与执行流程ONNX Runtime 提供了现成的quantize_dynamicAPI直接针对之前导出的 encoder 和 decoder 分别量化就行from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( encoder.onnx, encoder_int8.onnx, weight_typeQuantType.QInt8, op_types_to_quantize[MatMul], ) quantize_dynamic( decoder.onnx, decoder_int8.onnx, weight_typeQuantType.QInt8, op_types_to_quantize[MatMul], )op_types_to_quantize我限定为MatMul。Transformer 里面最占计算量的就是各种投影矩阵的矩阵乘法把这些量化成 int8 就能拿到主要收益又不会像全算子量化那样引入不必要的精度风险。量化完成后把之前的sess_enc和sess_dec换成对应的_int8.onnx即可translate函数完全不用改。5.3 量化前后对比体积、速度和翻译质量我在一台普通 i5 CPU 的旧笔记本上做过一轮对比模型就放在本地单条句子长度在 15 到 35 个词之间。注意不同机器差异很大这里看相对趋势更有参考价值方案权重总体积单句平均延迟翻译质量PyTorch transformers约 315MB约 800ms基线ONNX FP32约 320MB约 420ms与基线基本一致ONNX int8 动态量化约 83MB约 260ms个别长句有轻微用词差异体积下降非常直观四个模型文件加起来从 300 多 MB 降到了 80MB 左右部署时不管是打进镜像还是热更新加载压力都小很多。延迟方面ONNX FP32 对比 PyTorch 其实已经有提升主要来自 ONNX Runtime 的图优化和算子融合int8 再进一步把矩阵乘法换成整数运算CPU 友好度明显提高。翻译质量的差异我拿 20 条句子做过简单人工对比。短句几乎没有任何可感知差异长句偶尔会出现选词不同比如“温度在二十度左右”变成“温度大约二十度”但没有出现语法崩坏、漏译整句这种严重问题。如果你要交付的是严肃的正式翻译稿件建议量化后自己跑一遍核心测试集。5.4 量化后翻车怎么排查量化不是一键点完就能安心上的我记录几个最容易踩的坑。某个长句输出乱码先回退到 FP32 ONNX 跑同一句如果 FP32 正常、int8 异常说明是量化精度损失集中在特定长度区间可以只对 decoder 量化、encoder 保留 FP32通常能缓解量化后输出突然全变成同一个 token基本可以确定是encoder_hidden_states那个输入轴在量化时被打断了检查一下量化日志里有没有把MatMul后的Reshape也一起量掉必要时改op_types_to_quantize空白字符和标点丢失这不是模型问题是 tokenizer 解码时的skip_special_tokens行为对比量化前后输出时注意统一解码参数别让解码差异混进量化评估。6. 部署到服务的最后一公里FastAPI 封装和并发细节6.1 用 FastAPI 包一个最小翻译接口模型转好、验证完肯定要接成服务。我这里用一个 FastAPI 的例子收尾不复杂但回显了部署时真正要注意的点。import threading from fastapi import FastAPI from pydantic import BaseModel app FastAPI() # 全局只初始化一次 sess_enc ort.InferenceSession(encoder_int8.onnx, providers[CPUExecutionProvider]) sess_dec ort.InferenceSession(decoder_int8.onnx, providers[CPUExecutionProvider]) lock threading.Lock() class TranslateRequest(BaseModel): text: str max_new_tokens: int 128 app.post(/translate) def translate_api(req: TranslateRequest): with lock: result translate(req.text, max_new_tokensreq.max_new_tokens) return {translation: result}lock是我故意加上的。虽然实测同一份InferenceSession在多线程下直接调用大概率没问题但不同版本的 ONNX Runtime 行为不完全一样与其排查偶发崩溃不如直接加一把锁。翻译接口不是高吞吐场景锁带来的性能损耗微乎其微。如果你真的需要高并发正确做法是维护一个 session 池每个工作线程持有独立 session。6.2 预热和线上配置的几个细节服务启动后不要等第一个真实请求才编译图形。在 FastAPI 的startup事件里主动跑一次翻译把 ONNX Runtime 的图优化流程触发掉否则第一个用户会撞上几百毫秒甚至几秒的初始化延迟。另外接口层要对输入文本长度做限制。这个 ONNX 模型虽然支持动态序列长度但太长的输入会让 decoder 的生成时长接近失控还会让内存占用膨胀。我在接口里加了一个请求体长度校验超过两三百字就直接返回 400。翻译任务一次性处理太长文本本身也没有意义不如拆成摘要或段落级任务。6.3 和原 Transformers 部署方案的整体对比部署这步走完可以重新审视整个迁移的价值。原来容器里需要transformers torch sentencepiecetorch 全家桶上来就是几百 MB 起步服务进程内存占用也偏高。换成onnxruntime fastapi之后运行依赖收敛到一个推理库镜像体积、内存占用、冷启动时间全都下降了。而且因为宿主环境不再依赖 PyTorch后续把模型推到更边缘的设备上也顺理成章。比如之前有朋友问我能不能再往下压把一个 ONNX int8 的翻译模型跑到离线盒子上去那基本就是拿onnxruntime的轻量版本跑同样的图部分设备甚至需要再转成kmodel之类的格式量化过程中积累的经验还能复用。我个人的实际体会是把 HF 模型迁移到 ONNX真正花费时间的往往不是“导出”那一步而是把自回归生成循环重新组织到引擎外面的过程。只要理解了生成逻辑和模型结构工具本身反而不是瓶颈。这次只做到了动态量化和 FastAPI 服务化如果后面还有精力我会继续把 KV Cache 和 beam search 都补上让翻译质量和延迟再各进一步。
阅读完成 · 觉得有帮助？
)
