从 PyTorch 到 ONNX：手把手导出 Laya 决策模型并验证 1e-5 的精度对齐-平芜编程栈

【免费下载链接】laya

Run Laya, the open-source Jev-compatible System-1 decision model, from Node.js / TypeScript via ONNX Runtime

项目地址：https://gitcode.com/gh_mirrors/laya/laya

Laya（package.json 中的@receptron/laya）是一个开源的、与 Jev 兼容的System-1 决策模型：给它一个状态（工单、邮件、JSON 对象）和若干类型化问题，它在一次前向传播中返回每个问题的校准概率，而不是生成文本。本文带你从 PyTorch 检查点出发，用一行脚本把它导出为 ONNX 决策模型，并验证与 Python 参考实现的精度对齐误差约1e-5，最终在 Node.js / TypeScript 环境下无需 Python 即可推理。

1️⃣ Laya 是什么：不写文本，只做决策

Laya 由 ModernBERT-large 编码器 + 专用决策头组成（英文检查点约 421M 参数），回答三类问题：

问题类型	含义	返回
`choice`	从若干选项中选一个	每个选项的概率
`score`	有序量尺上的期望等级	等级分布 + 期望分数
`noul`	是/否陈述	校准的 P(true)

它的关键设计是运行时零 Python 依赖：src/laya.ts 基于onnxruntime-node加载 ONNX 图，所有推理都在 Node.js 进程内完成（CPU 上 3 个问题的批量调用约 140 ms）。

2️⃣ 为什么需要自己导出 ONNX

日常使用时Laya.load()会自动从 Hugging Face 下载现成的 ONNX 权重包（fp32，约 1.7 GB）并缓存到~/.cache/receptron-laya。只有在两种情况下需要自己动手导出：

你要基于更新版本的 PyTorch 检查点构建模型包；
你在使用一个尚未发布的模型变体。

仓库已内置完整工具链：导出脚本 export/export_onnx.py，模型包说明见 onnx/README.md。

3️⃣ 导出三步走：环境 → 权重 → 推理图

以下命令与 README.md 中 "Exporting the ONNX bundle yourself" 一节完全一致。

第一步：准备 Python 虚拟环境

cd export uv venv -p 3.12 .venv uv pip install -p .venv/bin/python torch transformers safetensors onnx onnxscript onnxruntime huggingface_hub

第二步：下载 PyTorch 检查点

只拉取导出所需的文件（权重、编码器、tokenizer、配置与参考实现代码）：

.venv/bin/python -c "from huggingface_hub import snapshot_download; snapshot_download('convaiinnovations/laya', local_dir='model', allow_patterns=['model.safetensors','encoder/*','tokenizer/*','rl_agent_config.json','rl_common.py','rl_agent_api.py'])"

第三步：运行导出脚本

.venv/bin/python export_onnx.py model ../onnx

脚本结束时会在终端打印**与 PyTorch 的最大 logit 差值（≈1e-5）**以及输出大小，这就是精度对齐的"体检报告"。

4️⃣ 脚本内部发生了什么

理解 export_onnx.py 的四个动作，你就理解了整个 ONNX 导出流程：

加载权重：从rl_agent_config.json构建模型（L24-L27），用safetensors严格加载权重并切换到 eval 模式。
动态形状导出：用torch.onnx.export（opset 18、dynamo=True）把 batch、序列长、选项数都声明为动态维度（L54-L62），导出的图对任意批大小通用。
附带 tokenizer 与校准参数：复制tokenizer/目录，并把max_len、head_max_len、按选项数分桶的 temperature 写入laya_config.json（L64-L67）——这些温度值是概率校准的关键，缺了它输出概率会失真。
精度对齐自检（parity check）：见下一节。

⚠️一个隐蔽的坑：脚本注释特别指出grad必须保持启用（L53）——PyTorch 的nn.TransformerEncoderLayer只在no_grad下走融合"fast path"，而那条路径不可导出。如果你在自定义脚本里包了torch.no_grad()再导出，很可能会失败。

导出的图有 5 个输入（input_ids、attention_mask、marker_pos、marker_mask、qtype）和 2 个输出（未校准logits、act_probs），接口定义见 onnx/README.md。

5️⃣ 1e-5 精度对齐是怎么验证的

导出的最后一步就是验证。脚本在 L69-L77 做了差分测试：

同一组样例张量分别跑PyTorch 参考实现和ONNX Runtime（CPU provider）；
对两组输出逐元素求差，打印max |dlogits|与max |dact|。

结果约为1e-5，含义是：ONNX 图与 PyTorch 的浮点偏差在第 5 位小数才出现。而 Laya 的最终输出概率按四位小数舍入（见 src/laya.ts 的round4），所以精度对齐误差远小于输出舍入粒度——对业务来说两者完全等价。

6️⃣ 在 Node.js 中加载并复核

把导出目录指向本地包，即可跳过下载直接使用：

const laya = await Laya.load({ modelDir: "./onnx" });

仓库提供了两层现成验证：

数值一致性测试：test/test_model.ts 用固定的"退款工单"样例，断言概率、分数、token 数与 Python 参考实现的输出逐位一致（如billing: 0.9415、urgency: 1.3886、churn_risk: 0.0988）。本地有模型包时运行LAYA_MODEL_DIR=./onnx yarn test即可执行。
可运行示例：examples/basic.ts 演示完整的systemOne调用与耗时统计，运行LAYA_MODEL_DIR=./onnx yarn example。

7️⃣ 常见问题速查

问题	说明
选项报 "options do not fit" 错误	每题选项须装进`head_max_len`（192 tokens），官方建议单个 choice 题少于 20 个选项
输入被截断	state 在问题头之后截断到`max_len`（英文检查点为 512 tokens），参数见 onnx/laya_config.json
内存不够	模型加载后约需 2 GB RAM，每批问题另加几百 MB
JS 与 Python 输出有微小差异	JSON 序列化需与 Python 的`json.dumps(ensure_ascii=False)`对齐，src/sequence.ts 的`pyJsonDumps`已处理，相关用例在 test/test_sequence.ts

8️⃣ 相关路径索引

导出脚本：export/export_onnx.py
模型包说明：onnx/README.md
校准参数：onnx/laya_config.json
运行时实现：src/laya.ts、src/sequence.ts
一致性测试：test/test_model.ts
调用示例：examples/basic.ts

总结：整条链路只有三行核心命令——建环境、下权重、跑脚本。脚本内置的 parity check 会当场告诉你精度对齐是否达标（1e-5 量级），之后用modelDir指向本地包，配合yarn test即可在纯 Node.js 环境中放心使用这份从 PyTorch 迁移而来的 Laya 决策模型。