第3章模型导出与 ONNX 转换

第3章：模型导出与 ONNX 转换

从 HuggingFace 导出模型，再到 ONNX，最后转成 TensorRT 引擎——这条链路我走了不下几十遍。说实话，每一步都有坑。今天我把踩过的雷、淌过的水，都给你捋一遍。

3.1 从 HuggingFace 导出模型：没那么简单

很多人觉得，从 HuggingFace 导模型不就是 model.save_pretrained() 吗？嗯，对了一半。但你要部署到生产环境，光保存权重是不够的。

我个人的习惯是：先确认模型类型，再决定导出策略。比如 LLaMA、ChatGLM、Qwen，它们的架构细节不同，导出方式也有差异。

核心原则：导出时保留完整的模型配置，包括 tokenizer、config.json、generation_config.json。缺一个，推理时就可能报错。

举个例子，我曾经部署一个 ChatGLM2 模型，导出时忘了保存 tokenizer。结果推理时 tokenizer 加载失败，整个服务起不来。排查了半天，最后发现是导出脚本少写了一行。

# 正确的导出姿势
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)

# 保存完整
model.save_pretrained("./chatglm2_export")
tokenizer.save_pretrained("./chatglm2_export")

小技巧：导出后检查一下目录结构，确认 config.json 和 tokenizer_config.json 都在。我每次都会跑一个 ls -la 看一眼。

3.2 ONNX 转换中的算子兼容性问题

ONNX 转换，说白了就是把 PyTorch 的计算图翻译成 ONNX 的中间表示。但翻译过程中，有些算子会「水土不服」。

最常见的坑：

动态控制流：PyTorch 里用 if 语句控制分支，ONNX 不支持。你想想看，ONNX 是静态图，它怎么知道运行时走哪条路？
自定义算子：HuggingFace 模型里经常有自定义的 Flash Attention、Fused LayerNorm 等。这些算子 ONNX 原生不认识。
数据类型不匹配：PyTorch 默认 float32，ONNX 可能转成 float16，精度对不上。

警告：遇到算子不兼容时，不要硬转。我见过有人强行转 ONNX，结果推理结果全是 NaN。正确的做法是：要么替换算子，要么用 ONNX Runtime 的扩展库。

我记得有一次，转一个 LLaMA 模型时，torch.where 算子死活转不过去。后来发现是 ONNX opset 版本太低，升级到 17 就解决了。

# 指定 opset 版本，避免算子兼容问题
import torch

dummy_input = torch.randint(0, 100, (1, 128))
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    opset_version=17,  # 推荐 17 或更高
    input_names=["input_ids"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch_size", 1: "seq_len"},
        "logits": {0: "batch_size", 1: "seq_len"}
    }
)

3.3 动态 shape 处理：绕不开的坎

生产环境里，输入的长度不可能固定。今天用户问一句「你好」，明天问一段「请写一篇 5000 字的文章」。所以动态 shape 是必须支持的。

但动态 shape 在 ONNX 里是个麻烦事。ONNX 默认是静态 shape，你要显式告诉它哪些维度是可变的。

上面代码里的 dynamic_axes 参数就是干这个的。我把它拆开讲一下：

0: "batch_size"：batch 维度可变，比如一次推理 1 条还是 8 条数据
1: "seq_len"：序列长度可变，比如 128 还是 2048

注意：动态 shape 会带来性能损失。因为 TensorRT 要为每种可能的 shape 都生成优化策略。我的建议是：如果业务场景里输入长度相对固定（比如 512 以内），就设成静态 shape，性能更好。

我曾经踩过一个雷：设了动态 shape 后，TensorRT 构建引擎花了 3 个小时。后来发现是 shape 范围设得太宽（1 到 8192），优化空间太大。改成 1 到 2048 后，构建时间降到了 20 分钟。

3.4 我踩过的 ONNX 转 TensorRT 的雷

这部分是我最想分享的。每一个雷，都是真金白银换来的教训。

雷区一：精度对齐问题

ONNX 转 TensorRT 后，推理结果和 PyTorch 对不上。这是最常见的问题。原因通常是：

TensorRT 默认用 float16，PyTorch 是 float32
某些算子（比如 LayerNorm）在 TensorRT 里有不同的实现

我的解决办法：先用 torch.allclose() 对比输出，误差超过 1e-3 就要排查。如果精度要求高，强制用 float32 推理。

# 强制使用 float32
import tensorrt as trt

builder = trt.Builder(logger)
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16)  # 默认开启 FP16
# 如果精度要求高，注释掉上面这行

雷区二：动态 shape 导致的内存爆炸

动态 shape 时，TensorRT 会为每个可能的 shape 分配内存。如果范围太大，显存直接爆掉。

我曾经部署一个 70B 模型，设了动态 shape 范围 1-4096。结果构建引擎时，显存占用飙到 80GB，直接 OOM。后来我把范围缩小到 1-1024，才跑通。

建议：动态 shape 的范围，根据业务实际需求来定。不要贪大。比如你的业务里最长输入是 1024，就设成 1-1024，别设成 1-8192。

雷区三：ONNX 图优化不足

ONNX 模型里可能有很多冗余节点，比如多余的 reshape、transpose。这些节点在 TensorRT 里会降低性能。

我的做法：先用 ONNX Simplifier 做一次图优化，再转 TensorRT。

# 安装 ONNX Simplifier
# pip install onnx-simplifier

import onnx
from onnxsim import simplify

model = onnx.load("model.onnx")
model_simp, check = simplify(model)
onnx.save(model_simp, "model_simplified.onnx")

雷区四：序列长度对齐问题

TensorRT 要求输入序列长度是 32 的倍数（某些架构是 64）。如果你的输入长度是 127，它会自动 padding 到 128。但 padding 会影响 attention 的计算。

解决办法：在 ONNX 导出时，显式指定 attention_mask，让模型知道哪些位置是 padding。

# 导出时带上 attention_mask
torch.onnx.export(
    model,
    (dummy_input, dummy_mask),  # 传入 input_ids 和 attention_mask
    "model.onnx",
    input_names=["input_ids", "attention_mask"],
    dynamic_axes={
        "input_ids": {0: "batch_size", 1: "seq_len"},
        "attention_mask": {0: "batch_size", 1: "seq_len"}
    }
)

3.5 本章知识体系

下面这张图，是我自己总结的 ONNX 转 TensorRT 的完整流程。你照着走，能少踩 80% 的坑。

这张图里，我特别标注了「常见问题」区域。你每次转模型时，对照着检查一遍，能省不少时间。

3.6 小结

模型导出和 ONNX 转换，是整个部署流程里最容易被忽视的一环。很多人觉得「能跑就行」，结果到了生产环境各种翻车。

我个人总结了三句话：

导出要全：权重、配置、tokenizer，一个都不能少
转换要稳：opset 版本、动态 shape、算子兼容，提前排查
验证要狠：精度对比、性能测试、边界情况，全部跑一遍

做到这三点，ONNX 转 TensorRT 的坑，你至少能避开 90%。剩下的 10%，嗯，那就要靠实战积累了。

📚 课程目录