[腾讯犀牛鸟2026] pnnx for torch exported program

[Edwardssss]

代码仓库：https://github.com/Edwardssss/ncnn （分支 pnnx-torch-exported-program）

上游：https://github.com/Tencent/ncnn

本项目为 pnnx 补齐了 exported program (pt2) 的完整转换支持，并完成了 transformers 4.x / 5.x 双版本兼容与多平台 CI 测试，已经提交了PR#6941

完成情况

项目要求	完成情况
pnnx 支持 torch dynamo exported program 直接转换（pnnx model.pt2）	新增 pt2 前端：load_exportedprogram.cpp（subgraph 翻译 + 入口）+ load_pt2_serde/tensor/defaults + load_exportedprogram_legacy（<2.8 容器）；消费 torch.export.save 产物；用法 pnnx model.pt2 inputshape=[1,3,224,224]
为 pnnx 单元测试添加 pt2 格式支持	顶层 398 个测试全部双路径（ts + pt2）与 eager 对比；导出失败自动 skip，并用 pt2_expectations 期望表 + 诊断子串做防静默审计；pt2 通道不设 skip 表（残留未转换的 ATen op 即硬失败，只剩 3 个 torch.export 自身的拒绝被钉住诊断）；另加 loader 白盒鲁棒测试（9 用例：负向护栏 + deflate 正向 + hostile-inflate 护栏）
解析 pt2 中的模型图和权重	JSON 图 + constants/weights 原始张量载荷；shape / strides / storage_offset / 共享存储视图 / 零 stride 展开 / 越界与超大 shape 护栏
默认参数恢复（Dispatcher 化）	append_default_kwargs 原 ~850 行 if-else 重构为 60 个 handler + defaults_dispatch 分派表，按 op 补 dim/dtype/keepdim 等默认并 reorder 到 pattern 顺序
legacy(<2.8) 容器支持	独立 legacy 读取模块（memzip + 内存嵌套 zip + pickle 状态字典），复用同一 tensor_meta 物化路径；6 个老 torch fixture 双版本验证
不引入第三方库	zip 解析复用自带 StoreZipReader（stored + deflate 方法-8，零依赖 RFC1951 inflate），JSON 用自带 pnnx_json；新增仅为解析逻辑并复用既有 pass 管线，零第三方依赖
通过多平台编译和测试	pnnx-pt2.yml ubuntu / windows / macos 全量 398 全绿；transformers 4.48.3 37 全绿（另备 5.x/4.x 双线 pre-push runner）；ncnn 后端回归 246/246（含 Tensor.new_empty / F.instance_norm 降级修复）；本地 Windows 395/395；本地全量 ctest 644/644（ts + pt2）

总体方案

torch.export(model, args)
        │ torch.export.save
        ▼
model.pt2 (exported program, zip 格式)
        │ pnnx load_exportedprogram
        ▼
pnnx 内部图 (aten op + 常量参数)
        │ pass_level0~5 + pass_ncnn
        ▼
pnnx 模型 (.pnnx.param/.bin) 或 ncnn 模型

前端按职责拆成多个模块：load_exportedprogram.cpp 负责 subgraph 翻译与入口；load_pt2_serde 集中 serde dtype / target 归一化；load_pt2_tensor 负责 tensor_meta → Attribute 物化（ts / pt2 / legacy 三路径共用）；load_pt2_defaults 用 handler Dispatcher（defaults_dispatch 分派表）为 dim/dtype/keepdim 等可选参数补默认值并按 pnnx pattern 顺序 reorder。与 TorchScript 路径共用 pass_level2/5 与 pass_ncnn 的算子重写逻辑，因此大部分算子可以零成本复用。

使用示例

导出并转换

import torch

class Model(torch.nn.Module):
    def forward(self, x):
        return torch.nn.functional.relu(x)

net = Model().eval()
x = torch.rand(1, 3, 224, 224)

program = torch.export.export(net, (x,))
torch.export.save(program, "model.pt2")

# 与 torchscript 路径相同的命令行
pnnx model.pt2
pnnx model.pt2 inputshape=[1,3,224,224] fp16=1
# 显式指定各输出文件
pnnx model.pt2 pnnxparam=model.pnnx.param pnnxbin=model.pnnx.bin \
       ncnnparam=model.ncnn.param ncnnbin=model.ncnn.bin

生成的产物与运行

转换后生成 model_pnnx.py、model.pnnx.param/.bin（ncnn 后端还有 model_ncnn.py 与 .ncnn.param/.bin）。生成的 python 模型自带确定性的 test_inference()（内部 torch.manual_seed(0)，与 eager 对齐）：

import model_pnnx
out = model_pnnx.test_inference()

在单元测试中走 pt2 路径

from pnnx_test_helper import test_pnnx
ok = test_pnnx(net, (x, y), ["[1,3,8,8]", "[1,3,8,8]"], "test_F_softmax")
# True  = pt2 转换正确
# False = 数值不匹配
# None  = 跳过（torch.export 不支持该模型）

关键点

1. exported program 的加载与符号内联

.pt2 是 zip 格式，参数与常量存储在 constants 中，图节点为 torch.fx.GraphModule 序列化。加载后，动态 shape 相关的 unbacked symbol（as_sym_int/as_sym_float/as_sym_bool）需要用具体输入求值并内联回图中，否则 pnnx 的 shape 推理会失败。

# pnnx_test_helper.py 中的 _inline_exported_symbols：
# 用 fx.Interpreter 跑一遍具体输入，捕获每个节点的符号参数值，
# 重写 pt2 zip 中的 JSON 图数据（as_sym_* → as_*），再回写 zip。

注意：pnnx 的 zip 读取器已支持 stored（不压缩）与 deflate（方法-8，RFC1951 raw-deflate）两类条目。重写 pt2 时仍建议用 zipfile.ZIP_STORED（与 torch.export.save 官方布局一致、兼容旧版本 pnnx）。

2. weight_norm 的语义修正

新版 torch.nn.utils.parametrizations.weight_norm 的范数是沿除 dim 外的所有维（norm_except_dim），与旧版 torch.nn.utils.weight_norm（沿 dim）不同。pnnx 最初按旧语义折叠，导致 Linear/Conv 的权重算错。

• 修复 torch_weight_norm_fold.cpp：把常量 v、g 折叠为最终权重，范数索引 = coord[dim]（沿除 dim 外求和），并处理 dim=-1（标量范数）与标量 g。
• GraphRewriter 的 replace_pattern_graph() 必须声明全部 pattern 输入（含未使用的 g），且 Attribute 输出名要显式写出，否则 insert-multiple 分支会空指针崩溃。
• match() 里提取 Attribute 数据缓存到成员（write 阶段匹配节点已被删除，不能再用指针访问）。

3. transformers 5.x 兼容

transformers 5.x 移除了 CLIPAttention 的 causal_attention_mask / output_attentions 参数，Qwen2Attention 引入 MLA（q_lora_rank/kv_lora_rank）。

• 用 _is_tf5 = version.parse(_transformers_version) >= version.parse('5.0') 分叉 forward 与输入生成；
• 4.x 不存在的参数（MLA）用 guard 让测试在 4.x 下 return True 跳过；
• 4.x/5.x 的随机输入生成顺序必须与 pnnx 生成的 test_inference()（torch.manual_seed(0) 序列）严格一致，否则比较必然失败（clip 的 casual_mask0 曾因生成顺序错误导致 4.x 全失败）。

4. legacy(<2.8) 容器支持

torch.export.save 从 torch 2.8 起才使用当前 zip 布局；更早版本（2.5..2.7）保存的是 serialized_exported_program.json + 序列化 state_dict/constants 的扁平归档。为此新增独立 legacy 读取模块（load_exportedprogram_legacy）：内存里展开嵌套 zip、用自带 pickle 协议-2 VM 解码 _rebuild_tensor_v2/_rebuild_parameter 并把 storage 分片还原为与 2.8+ 一致的 fqn → (payload, tensor_meta) 映射，随后复用同一 load_pt2_tensor 物化路径。测试用 torch（2.7.1）生成 6 个 fixture，在当前 torch 下做双版本回归。

5. loader 的模块化重构

原文件 load_exportedprogram.cpp按功能拆分并做了两处结构性重构：

• serde dtype/归一化、tensor_meta 物化、默认参数补全各自独立成模块，主文件只剩 subgraph 翻译与入口；
• 默认参数 Dispatcher：append_default_kwargs 的巨型 if-else 改为 60 个 handler + defaults_dispatch 分派表，加新 op 只需一个 handler + 一条表项；
• 鲁棒性白盒化（负向 + 正向）：test_pt2_robustness 把 loader 的防御护栏固化为负向用例（缺权重 record、超大 shape / storage_offset 越界、截断 zip、恶意 deflate 头等），并把 deflate 压缩样本固化为正向用例（须解压无损、产物与 stored 逐字节一致）；pt2 期望表用诊断子串做防静默审计；
• deflate 条目支持：StoreZipReader 放行 method-8（raw deflate）并新增零依赖 RFC1951 inflate（stored / fixed-huffman / dynamic-huffman 三类 block），read_file 按需解压、get_file_size 返回解压后大小；官方 .pt2（stored）行为不变；
• ncnn 通道补强：convert_aten_new_empty 支持 ts 侧 Tensor.new_empty（deepseek MLA kv-cache 缓冲），F_instance_norm pattern 补 use_input_stats，ncnn 后端测试 246 项全部通过。

6. torchaudio spectrogram / iSTFT：把 pt2 的展开链整体降到 ncnn

torch.export 会把 torchaudio.functional.spectrogram 内联展开成 stft + view_as_real（view_as_complex）+ abs( square ) / div( sqrt ) / sum 一串 aten op，而 pass_ncnn 原有的 torch_stft 模式只匹配「单个 functional op」，pattern 永远不命中，最终图上剩一堆未转换的 ATen op —— pt2 链路失败。

于是我们在 pass_ncnn 新增 pt2 变体（torch_stft_pt2_{pad_complex,complex,power,abs,norm}、torch_istft_pt2{,_norm}），把整条展开链直接降到 ncnn Spectrogram / InverseSpectrogram，并在 match() 里进行逐条验证。能覆盖到power=None/1/2、normalized='window' / 'frame_length'、显式 pad、constant / reflect 两种 pad；前向模型从 38 层降到 19 层（5 个 Spectrogram），逆向从 17 降到 8（4 个 InverseSpectrogram），0 个 ATen 残留。

复现校验：用 3 个不同版本的 torch（2.11.0 / 2.12.1 / 2.13.0）导出的 pt2 图都能降到 5 + 4 个真实层、0 残留。

多平台 CI

新增 .github/workflows/pnnx-pt2.yml：

Job	平台	transformers	范围
pt2	ubuntu / windows / macos	5.16.1	顶层 398 项测试
pt2-transformers4	ubuntu	4.48.3	transformers 37 项测试

顶层 transformers 测试只需 pnnx 二进制；依赖全部固定（torch 2.13.0 / torchvision 0.28.0 / torchaudio 2.11.0 / transformers 5.16.1 / 4.48.3）。

关于 transformers 4.x 的 CI测试

pt2-transformers4 是向后兼容的回归保护测试。我们的 test_transformers_*.py 为适配 transformers 5.x（其 CLIPAttention / nn.MultiheadAttention / MLA 接口大改）引入了 _is_tf5 分支，但 4.x 仍是主流稳定版，为 5.x 做的改动很容易悄悄破坏 4.x 路径。本地用 tf4env（transformers 4.48.3）验证时，clip 的 casual_mask0 随机输入生成顺序与 pnnx 生成的 test_inference() 不一致，导致测试失败。因此单独跑一个 ubuntu + 37 项的测试，可以用于确认为 5.x 做的适配没有破坏 4.x 兼容。

关于 Windows runner 上的 MKL 崩溃

Windows CI 曾偶发失败于 test_torch_fft_{hfftn,ihfftn,irfftn}：TorchScript 路径在 pass_level0（shape_inference 实际执行 jit 图）阶段静默崩溃，导致 _pnnx.py 未生成、测试以 ModuleNotFoundError 失败。但是我在本地Windows环境下能够成功，最初我仔细排查了import路径拼接，环境变量等问题，均未能成功复现。

根因排查：

• 失败只发生在 GitHub-hosted Windows Server 2025 runner，本地 Win10 及 ubuntu/macos 均无法复现（本地尝试复刻 CI 循环 100 轮测试全部通过）；
• 三个失败测试都是 half-spectrum 多维 FFT（hfftn/ihfftn/irfftn），torch 2.13 CPU 版的 MKL FFT 在 CI CPU 上分派到新指令集 kernel 时对这类形状崩溃；

修复：在 pnnx-pt2.yml 的两个测试步骤中加上限制：

export OMP_THREAD_LIMIT=1
export OMP_NUM_THREADS=1
export MKL_NUM_THREADS=1
export MKL_ENABLE_INSTRUCTIONS=SSE4_2

验证结果

验证项	结果
顶层全量（398，ts+pt2，含 legacy + loader 鲁棒白盒）	398 / 0 失败
loader 白盒鲁棒 test_pt2_robustness（负向 + deflate 正向 + hostile-inflate 护栏）	9 / 0 失败
transformers 5.16.1（37）	37 / 0 失败
transformers 4.48.3（37）	37 / 0 失败
ncnn 后端回归（246）	246 / 0 失败
CI pt2：ubuntu / windows / macos（398）	398 / 0 失败
CI pt2-transformers4：ubuntu（37）	37 / 0 失败
torchaudio spectrogram / iSTFT 家族（10，含 power=1 变体）	10 / 0 失败
折叠类（full / ones / zeros / new_full / arange，13）	13 / 0 失败
多 torch 生产者 pt2 降低（2.11.0 / 2.12.1 / 2.13.0）	5 + 4 层、0 残留

已知限制

• 依赖动态控制流/自定义 autograd 的图仍回退到 TorchScript 路径。没有 pnnx 侧的 skip 表：pt2 图上残留未转换的 ATen op 就是硬失败；
• 静态但超大的常量折叠（> 1 GiB）会在 match() 阶段被拒绝，转换在 ncnn 边界显式失败，而不是产出一个 shape 与字节数不一致的坏常量；
• spectrogram 的少数组合仍未覆盖（power=2 + 窗口归一化、power=None + 任意归一化）；
• Qwen2/Qwen3/DeepSeek 的 MLA 等算子走 ncnn 自定义层或 pass 展开，未来可继续下沉；
• 动态 shape 需显式传 inputshape（symbolic 维度由输入 shape 覆盖后固定化转换），未保留运行时的动态维度断言。