跳到主要内容

从 PyTorch 到 Torq™ NPU：将 AI 模型转换为高效部署的二进制

· 阅读需 12 分钟
Ye Htet
Ye Htet
Embedded AI @ Synaptics

这是一篇与实时语音识别博客配套的文章。后者专门介绍了 Moonshine V2 的导出过程。这里我想从更高层面展开，说明 torq-tools 用来把任意 PyTorch 模型转成可在 Torq T1 NPU 上运行的 .vmfb 二进制文件的通用方法。如果你有自己的模型希望移植到开发板上，这就是需要遵循的整体工作流。

引言​

PyTorch 已成为开发和训练机器学习模型的标准框架，因此它也成为大多数 AI 部署的起点。虽然 PyTorch 模型很容易训练和评估，但它不能直接在 NPU 上运行。就像优化过的 C++ 代码一样，神经网络也必须被转换、优化并编译成目标硬件能够高效执行的格式。本文将介绍如何把 PyTorch 模型转换成一个可用于 Torq T1 NPU 上高性能推理的编译后二进制文件，并以 Astra SL2610 内的 Torq T1 为目标。

为什么要在 NPU 上运行模型？​

像 Torq NPU 这样的 NPU 是为 AI 模型运行而专门设计的。它们提供更高的性能和效率，并让 CPU 能够继续处理其他任务。

总览​

图示展示了从 PyTorch 模型到可在 Torq NPU 上运行的编译后 Virtual Machine FlatBuffer（.vmfb）二进制文件的工作流。

workflow

模型的起始组件​

PyTorch 和 Hugging Face 的 Transformers 是训练 AI 模型的主流框架，因为它们消除了从想法到可运行模型的大部分摩擦。PyTorch 的动态、Python 化设计让实验和调试变得简单。Transformers 则在这个生态上构建了更高层的抽象，为开发者提供了对数千个预训练模型的直接访问，以及简单的微调 API，从而无需从头训练。

ONNX（Open Neural Network Exchange）是一种用于表示训练后模型的框架无关静态计算图的格式。由于 PyTorch、TensorFlow 等大多数框架都能导出到 ONNX，因此它已成为编译器工具链的常见起点。

为 NPU 优化模型​

Torq NPU 被设计成处理静态形状，因此模型中任何动态形状都需要先被固定为静态。也就是说，要么从 PyTorch 导出时就指定固定形状，要么在导出后对所有动态维度做图修补（graph surgery）。

量化模型参数——例如将 32 位浮点数转换成 8 位整数——也能将内存占用大约减少 4 倍，并显著加速推理，因为低精度运算在 NPU 上更快，并且传输的数据更少。

为什么一开始就从 PyTorch 开始？​

有时你可以直接在 Hugging Face 或其他地方找到某个模型的 ONNX 导出结果。虽然 ONNX 文件可以被快速转换并编译到 Torq NPU 上，但请记住：ONNX 图本身是冻结的。如果其中存在动态形状或不受支持的算子，后续就会遇到问题。任何不受支持的 ONNX 模型部分都只能靠后处理 workaround，而不是在源头修正——而从 PyTorch 直接出发则能避免这个问题。

在 PyTorch 层面工作，从 model.safetensors 权重和建模代码出发，你在图被冻结前仍有机会进行修改：重写模型的前向传播使其更适合导出、将单体模型拆分为多个组件，或替换成自定义的 attention 实现。一旦模型变成 ONNX，这些修改会变得困难得多（此时你需要做的是图修补，而不是直接改 Python）。不过，尽管如此，有些场景还是必须以 ONNX 作为起点。

引入 torq-tools​

Synaptics 提供了一组开源工具，名为 torq-tools，用于导出模型并让其在 Torq NPU 上运行。这些工具正是本文讨论的各类主题的实践实现。

请查看 torq-tools 的 README.md，了解如何安装和使用这些工具。

导出脚本​

torq-tools 包含多个流行模型的导出脚本，包括 Moonshine、Moonshine Streaming、SmolLM2、Gemma-3 和 LFM2.5。这个列表还会持续增长。

每个模型导出脚本都位于 src/torq/models/<model>/export.py 下，它们都是本文所描述通用流程的具体实例，并都继承自共享基类：src/torq/model_export/onnx.py 中的 OnnxModelExporterBase。

导出器的旅程​

模型导出器的典型流程通常是这样的：

PyTorch weights ──► ONNX (raw) ──► ONNX (static, edited) ──► ONNX (bf16/quantized) ──► MLIR ──► .vmfb

下面概述这个旅程中的每个阶段。

第 1 步：从 PyTorch 导出 ONNX 图​

常见路径有两种，不同的导出器会用到其中的一种：

  • torch.onnx.export(..., dynamo=True) —— 直接从模型的 nn.Module 导出。这是你需要在导出前手工修改模型时的选择（例如 Moonshine 拆分成 frontend / encoder / adapter / cross_kv / decoder_kv 组件）。
  • optimum-cli export onnx —— 当 Hugging Face 的 optimum 已经知道如何追踪某个模型时，可以用这个方法（SmolLM2、Gemma-3、LFM2.5 都通过 torq.model_export.hf.optimum_export_onnx 使用这一路径）。这段代码更少，但你会放弃在导出时精细控制图结构的能力——任何后续修改都只能通过图修补来完成。
# src/torq/model_export/hf.py
def optimum_export_onnx(onnx_dir, hf_repo, dtype, models, *, opset=22, opt_level="O1"):
    cmd = [
        sys.executable, "-m", "optimum.commands.optimum_cli", "export", "onnx",
        str(onnx_dir), "--model", hf_repo, "--dtype", dtype, "--opset", str(opset),
    ]
    ...

无论采用哪种方式，你最终都会得到一个或多个带有动态形状的 .onnx 文件（batch、sequence length、KV-cache length 都是符号维度）——这就是后续所有步骤的输入。

第 2 步：清理和规范化图​

在处理图之前，每个导出器都会先经过 onnx_graphsurgeon 处理，去掉无效结构并把图整理成规范形式：

graph = gs.import_onnx(model)
graph.name = "main"
graph.cleanup(
    remove_unused_graph_inputs=True, remove_unused_node_outputs=True
).toposort()
model = gs.export_onnx(graph)

命名也很重要：MLIR 标识符不接受 ONNX 允许的所有字符，因此 OnnxModelExporterBase.sanitize_onnx_names 会在模型真正保存前重写张量/初始化器名称，去掉所有不在 [a-zA-Z0-9_./] 范围内的字符。忽略这一步是后面出现神秘 MLIR 解析失败的常见原因。

第 3 步：让图变成静态图​

这是最需要按模型定制工作的一步，因为Torq NPU 需要固定长度的输入/输出和操作——没有符号/动态形状，也没有动态控制流。每个导出器都会实现一个抽象的 make_static() 方法，它负责：

  1. 修正 I/O 形状 —— batch size、sequence length、音频块长度、KV-cache 长度都变成写进图中的具体整数。
  2. 用固定大小缓冲区替换动态 KV-cache 增长 —— 自回归解码器通常会按每个 token 增长一格 KV cache；静态图则会提前分配完整缓存，并在计算出的偏移位置写入。
  3. 增加显式 causal/attention mask —— 由于图无法再从当前动态序列长度中推断掩码，必须根据固定长度显式添加 mask。
  4. 将动态索引计算转成静态索引计算 —— 例如 Range(start, start + 1, 1) 这样的模式（用于计算“下一次写入位置”）会被改写成不需要在 trace 阶段使用动态 range 的形式。

在 torq-tools 中，这些都被实现为可复用的图编辑 pass（ReplaceDynamicKVCache、MaskFutureAttentionScores、AddCurrLenInput、ConvertToStaticIndex，位于 src/torq/graph_edit/edits/transformer.py），并按具体模型组合起来：

# src/torq/models/smollm2/export.py (abbreviated)
(
    editor
    .replace_dynamic_kv_cache(cur_len, self._max_gen_tokens)
    .mask_future_attn_scores(cur_len, self._max_gen_tokens)
    .add_curr_len_input(cur_len)
    .convert_to_static_index()
)

完成这一步之后，export_onnx() 会显式检查是否还存在动态维度：

dynamic_shapes = check_dynamic_shapes(onnx.load(self._export_paths[comp]))
if dynamic_shapes:
    raise ValueError(f"Model '{comp}' still has dynamic shapes: {json.dumps(dynamic_shapes)}")

第 4 步：为编译器和硬件兼容性做图修补​

即使图已经完全静态，它仍可能包含 Torq 编译器暂不支持的算子，或者包含对 NPU 来说冗余或昂贵的形状/模式。src/torq/graph_edit/edits/ 按关注点组织：

模块处理内容
arithmetic.py分解 LayerNormalization，折叠标量 MatMul 到 Mul，去除冗余 cast / IsNaN 检查，替换常量 Div 为 Mul
conv.py分解 strided 1D convolution，扩宽 strided depthwise conv
shape.py消除无操作 Transpose/Expand，压缩 Reshape 链，广播 op 输入
transformer.py静态 KV-cache / attention-mask 机制（见第 3 步），压缩 GQA broadcast，重定向 cross-attention key layout
rnn.py将双向 RNN 分解为前向/后向两半
padding.py吸收 padding 到相邻 op，重写负 pad，替换 Pad 为 Concat
custom_ops.py将 fused ops 如 GroupQueryAttention / SimplifiedLayerNorm 替换为基础算子等价形式
artifacts.py抽取常量查找表（例如 token embeddings）到图外的 .npy 文件中，修剪/拆分 LM head

这也是我们分解当前编译器版本不支持的算子的地方：在 Moonshine 的例子里，asinh 需要用拟合多项式近似，而部分 LayerNorm 因 MLIR lowering 问题需要被分解。这里要保守处理——过度分解可能损害精度或延迟，因此只分解那些实际不受支持或用原始形式明显更慢的部分。

每个导出器都会在 apply_post_static_patches() 中应用自己的子集，且每个导出组件只调用一次。

第 5 步：在继续前进行校验​

每个导出器都会实现 validate_onnx()，它会在一组提示词/输入上把静态 ONNX 模型与原始动态模型并行运行，并比较输出：

output = runner.run(input)          # static model
val_output = val_runner.run(input)  # original dynamic model
if output[:min_len] != val_output[:min_len]:
    result = "Warning: Validation failed, mismatched outputs..."

在做 dtype 转换或量化之前执行这一步非常重要。这样一来，如果后面出问题，你已经知道静态 fp32 图本身是正确的。

第 6 步：精度转换和量化​

在此时，两个独立且可选的工具可以分别作用于静态 ONNX 模型：

Dtype conversion（torq.tools.convert_dtype）会把 fp32 → bf16（Torq 运行时原生支持 bf16），并在合法时把 int64 → int32/int16/int8：

torq-convert-dtype onnx -d bf16 -i model_fp32.onnx -o model_bf16.onnx
torq-convert-dtype onnx -d int32 -i model_bf16.onnx -o model_bf16_int32.onnx

Weight quantization（torq.tools.quantization.weight_quantization）则更进一步，把 MatMul 权重量化为 int8 或 int4。对于类 LLM 的解码器，analyze 步骤会基于校准 prompt 执行 KL 散度敏感性分析，并按层自动将每层分桶到合适的 bit-width：

torq-quantize-model analyze -i model_fp32.onnx -o sensitivity.json \
    --config-output quant_config.json --embeddings token_embeddings.npy --bits 4 8

torq-quantize-model quantize -i model_fp32.onnx -o model_mixed_bf16.onnx \
    --config quant_config.json --dequantize-weights

--dequantize-weights 会把量化误差直接写入一个普通 bf16 模型中（可直接用于编译）；如果不加它，你会得到 DequantizeLinear 节点，这更适合在最终定型前检查或分享量化权重。

第 7 步：ONNX → MLIR → VMFB​

这是每个模型都一致的最后一段，由 src/torq/utils/compile.py 处理。它是一个两跳 lowering：

  1. iree-import-onnx（或等价的 Python API iree.compiler.tools.import_onnx）将 ONNX 图转换为 IREE 风格的 MLIR。
  2. torq-compile 把该 MLIR 进一步降低到目标为 torq backend 的 .vmfb 二进制。
def export_torq(input_model, output_dir, ...):
    ...
    export_onnx_to_mlir(input_model, mlir_model, opset=opset)
    compile_mlir_for_vm(mlir_model, vmfb_model, compiler_args=compiler_args, ...)

如果你已经有一个静态 ONNX 文件，只想拿到二进制，也可以直接用这个单一 CLI：

torq-compile model_bf16.onnx -o model.vmfb

默认会交叉编译到开发板的 aarch64 目标；传入 --local-compile 可以在主机上直接编译并运行，用于快速迭代（当你手边没有开发板，或想更快排查编译错误时非常有用）。

把这些步骤串起来​

对于已有注册导出器的模型，上述 7 个步骤会作为一条命令执行：

torq-export-model <model_name> [model-specific args] --convert-dtypes

例如 Gemma-3：

torq-export-model gemma3 --model-size 270m --instruct-model --convert-dtypes

导出器会按顺序执行：下载 HF 模型 → 导出 ONNX → 使其静态化 → 应用图修补 → 校验 → 转到 bf16 → 编译成 .vmfb，最后得到如下目录结构：

models/<hf_repo>/export/onnx/float/static/      # static fp32 ONNX + validation
models/<hf_repo>/export/onnx/converted/static/  # bf16/int32-converted ONNX
models/<hf_repo>/export/torq/converted/static/  # .mlir + .vmfb

编写你自己的导出器​

如果你的模型不属于现有配方，你可以通过继承 OnnxModelExporterBase 来编写自己的导出器。接口很小——只需要实现：

方法责任
_setup_dirs()定义 source/export/converted/torq 产物的保存位置
_load_onnx()从 PyTorch/HF 获取每个组件的原始 ONNX ModelProto
make_static()Fix shapes, replace dynamic KV-cache/control-flow with static equivalents
apply_post_static_patches()Per-component graph edits (op decomposition, embedding extraction, etc.)
validate_onnx()Parity check vs. the original dynamic model

其他内容——清理名字、dtype 转换、量化、MLIR lowering、VMFB 编译——都是共享基础设施，开箱即用。

总结​

  • 把模型部署到 Torq NPU 的起点是 PyTorch，而不是现成的 ONNX 文件。这样你仍然保留在图冻结前修改模型的能力。
  • 图必须完全静态化：固定形状、固定大小 KV cache、显式 mask、静态索引计算。
  • 图修补用于修正编译器暂时无法 lower 的内容，按关注点组织（算术、卷积、形状、transformer、rnn、padding、自定义算子）。
  • 在进行精度转换之前，先验证静态 fp32 图是否与原始动态图一致。
  • 精度转换（bf16）和量化（int8/int4）是独立的、可选的、可组合的步骤。
  • ONNX → MLIR → VMFB 的降低流程对所有导出器都是通用的、共享的。