HY-MT1.5-1.8B-onnx

HY-MT1.5-1.8B ONNX

这是 Tencent HY-MT1.5-1.8B 翻译模型的 ONNX 导出版本，可直接使用 ONNX Runtime 进行高效推理，无需 PyTorch 依赖。

模型架构

导出为双模型架构，适配自回归文本生成的 KV Cache 推理模式：

输入输出规格

Prefill 模型（2 输入 → 65 输出）：

Decode 模型（66 输入 → 65 输出）：

快速开始

环境要求

pip install onnxruntime transformers==4.56.0

推理示例

import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer

# ── 加载 ──────────────────────────────────────────
tokenizer = AutoTokenizer.from_pretrained("tencent/HY-MT1.5-1.8B")

prefill = ort.InferenceSession("onnx/hy_mt_prefill.onnx",
                               providers=["CPUExecutionProvider"])
decode  = ort.InferenceSession("onnx/hy_mt_decode.onnx",
                               providers=["CPUExecutionProvider"])

# ── Tokenize ──────────────────────────────────────
messages = [{"role": "user", "content": "将下面的英文翻译成中文: Hello, how are you?"}]
input_ids = tokenizer.apply_chat_template(
    messages, tokenize=True, add_generation_prompt=True,
    return_tensors="np",
)

# ── Prefill ───────────────────────────────────────
seq_len = input_ids.shape[1]
attn_mask = np.ones((1, seq_len), dtype=np.int64)
prefill_out = prefill.run(None, {
    "input_ids": input_ids,
    "attention_mask": attn_mask,
})

# 提取 KV Cache
kv = {}
for i in range(32):
    kv[f"past_key_{i}"]   = prefill_out[1 + 2*i]
    kv[f"past_value_{i}"] = prefill_out[1 + 2*i + 1]

# 采样第一个 token
logits = prefill_out[0][0, -1, :]
next_token = int(np.argmax(logits))

# ── Decode 循环 ──────────────────────────────────
total_len = seq_len
generated = []
for _ in range(256):
    total_len += 1
    decode_out = decode.run(None, {
        "input_ids": np.array([[next_token]], dtype=np.int64),
        "attention_mask": np.ones((1, total_len), dtype=np.int64),
        **kv,
    })
    # 更新 KV Cache
    for i in range(32):
        kv[f"past_key_{i}"]   = decode_out[1 + 2*i]
        kv[f"past_value_{i}"] = decode_out[1 + 2*i + 1]

    next_token = int(np.argmax(decode_out[0][0, -1, :]))
    if next_token == 120020:  # EOS
        break
    generated.append(next_token)

# ── Detokenize ────────────────────────────────────
output = tokenizer.decode(generated, skip_special_tokens=True)
print(output)

更完整的推理脚本见 inference_onnx.py。

ONNX vs PyTorch 精度对比

使用 do_constant_folding=False 导出可保证 ONNX 与 PyTorch 输出完全一致：

导出说明

导出环境

关键配置

• 精度: float32（从 bf16 转换）
• 注意力实现: eager（sdpa 不可被 ONNX trace）
• 常量折叠: 关闭（do_constant_folding=False，否则会因 seq_len 被烘焙为常量导致动态输入时数值错误）
• 导出方式: TorchScript trace（dynamo export 与 HunYuan 自定义算子不兼容）

重新导出

python export_onnx_full.py

导出脚本会自动验证 ONNX 与 PyTorch 输出一致性。

性能参考

推荐使用 CUDAExecutionProvider 获得更好性能。

文件结构

.
├── onnx
├── export_onnx_full.py             # 完整导出脚本（prefill + decode）
├── inference_onnx.py               # ONNX Runtime 推理脚本
├── compare_onnx_pt.py              # ONNX vs PyTorch 精度对比
└── README.md

引用

@misc{hy-mt1.5,
      title={HY-MT1.5 Technical Report},
      author={Mao Zheng and Zheng Li and Tao Chen and Mingyang Song and Di Wang},
      year={2025},
      eprint={2512.24092},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2512.24092},
}