3. 模型格式概述：常见模型格式（PyTorch、ONNX、TensorRT）及转换需求

3.1 核心原理：三种格式的定位差异

在深度学习部署链路中，PyTorch、ONNX 与 TensorRT 分别承担着“训练研发”、“模型交换”与“生产推理”三种截然不同的角色。理解它们的本质差异，是掌握 TensorRT-LLM 转换流程的前提。

  • PyTorch（.pt / .pth）：面向动态图训练与实验。权重以 Python 对象形式保存，包含完整的计算图结构、优化器状态甚至源码级类定义。其优势是灵活，但缺点是运行时依赖 Python 解释器，且算子粒度偏细，不利于底层硬件优化。
  • ONNX（.onnx）：一种静态计算图的中间表示。它将模型描述为算子（Op）与张量（Tensor）的有向无环图，不绑定任何训练框架。ONNX 的价值在于“桥接”——让 PyTorch 模型可以导出为统一格式，再被 TensorRT、OpenVINO、ONNX Runtime 等推理引擎消费。
  • TensorRT（.engine / .plan）：NVIDIA 针对其 GPU 架构深度优化的推理引擎格式。它不再保留原始网络结构，而是将层融合、精度校准、内核自动调优后的执行计划（Execution Plan）固化下来。TensorRT 格式与具体 GPU 型号、驱动版本、TensorRT 版本强绑定。

从转换视角看，三者构成一条“由软到硬”的流水线：PyTorch 负责产出参数与结构，ONNX 负责标准化表达，TensorRT 负责将标准化图映射为高性能 GPU 指令。TensorRT-LLM 虽然也支持直接加载 HuggingFace 格式的权重，但内部仍会经历“解析→图优化→编译”的类似过程。

关键认知：不要将 ONNX 视为“最终部署格式”。ONNX 只是中间载体，真正的性能收益来自 TensorRT 的层融合（如将 QKV 投影合并）、内核自动调优以及 KV Cache 专用内存管理。
训练研发阶段 模型交换阶段 生产推理阶段 PyTorch .pt / .pth / .safetensors 动态图 · Python 运行时 权重 + 网络结构 + 优化器 适合训练 / 实验 / 断点续训 ⚠ 推理延迟高 · 显存占用大 ONNX .onnx 静态图 · 算子级 DAG 框架中立 · 可移植 适合模型交换 / 跨引擎迁移 ⚠ 无硬件优化 · 算子兼容风险 TensorRT .engine / .plan 层融合 · 内核自动调优 FP16/INT8 量化 · KV Cache GPU 专属执行计划 ⚠ 与 GPU 型号 / 版本强绑定 torch.onnx.export trtexec / onnx-tensorrt TensorRT-LLM 可直接加载 HF 权重，也可通过 ONNX 图转换，最终编译为 TensorRT engine

3.2 方法流程：PyTorch → ONNX → TensorRT 的转换链路

标准转换流程分为三个步骤，每一步都有独立的验证要点。以 Llama-7B 为例，典型命令如下：

# 第一步：PyTorch → ONNX（需固定输入维度，关闭动态控制流）
python convert_llama_to_onnx.py \
  --model_name meta-llama/Llama-2-7b-hf \
  --output llama7b.onnx \
  --seq_len 2048 \
  --dtype fp16

# 第二步：ONNX → TensorRT（使用 trtexec 或 Python API）
trtexec --onnx=llama7b.onnx \
  --saveEngine=llama7b.engine \
  --fp16 \
  --workspace=8192 \
  --minShapes=input_ids:1x1 \
  --optShapes=input_ids:1x2048 \
  --maxShapes=input_ids:8x2048

# 第三步：TensorRT-LLM 加载（推荐直接使用 HF 权重）
python convert_checkpoint.py \
  --model_dir meta-llama/Llama-2-7b-hf \
  --output_dir ./tllm_checkpoint \
  --dtype float16

实际工程中，TensorRT-LLM 官方推荐跳过 ONNX，直接利用其内置的 convert_checkpoint.py 脚本将 HuggingFace 权重转换为 TensorRT-LLM 的 checkpoint 格式，再通过 trtllm-build 生成 engine。原因在于：ONNX 导出大模型时容易遇到算子不支持、动态形状处理繁琐、图优化信息丢失等问题，而直接解析 HF 权重可以保留更多注意力结构细节（如 GQA、RoPE 等）。

转换路径 适用场景 优点 缺点
HF → TensorRT-LLM（推荐） LLM 推理部署 保留注意力结构、支持 GQA/RoPE、转换快 仅限 TensorRT-LLM 生态
PyTorch → ONNX → TensorRT 通用模型（CNN、非 LLM） 框架中立、可跨引擎复用 大模型导出易失败、动态形状处理复杂
PyTorch → TensorRT（torch2trt） 快速原型验证 无需中间格式 算子覆盖有限、维护不活跃

3.3 案例/场景：一个实际转换中的“坑”

某团队尝试将 ChatGLM2-6B 从 PyTorch 导出为 ONNX 再转 TensorRT，遇到了典型的 动态形状问题。原始模型在生成时，past_key_values 的长度会随解码步数增长，导致 ONNX 图中出现动态维度。使用 trtexec 转换时报错：

[E] Error[0]: The provided minimum shape is not compatible with the ONNX input shape.
Input: "input_ids" with shape: [-1, -1]

解决方案是采用 TensorRT-LLM 的 Inflight Batching 机制，它原生支持可变序列长度，无需固定到最大长度。团队最终放弃 ONNX 路径，直接使用 convert_checkpoint.py 加载 HF 权重，并设置 --max_seq_len 8192，成功生成 engine。这一案例说明：对于 LLM，ONNX 的静态图特性反而成为负担，TensorRT-LLM 的专用转换工具才是正解。

常见误区：不要认为“ONNX 是万能的中间格式”。LLM 中的 past_key_values、position_ids、attention_mask 等动态结构在 ONNX 中表达非常别扭，且 ONNX Runtime 与 TensorRT 对算子支持范围不一致，经常需要手工替换算子（如将 Gather 替换为 Slice）。

3.4 误区/复盘：格式转换中的五个典型错误

  1. 忽略版本绑定：TensorRT engine 与 CUDA、cuDNN、TensorRT 版本严格绑定。在一台机器上生成的 engine，换到另一台机器可能直接加载失败。务必使用 trtexec --dumpProfile 验证兼容性。
  2. 混淆权重格式与推理格式：.safetensors 只是权重存储格式，不包含计算图；而 .engine 是编译后的执行计划。两者不可直接互换。
  3. 忽视精度声明：PyTorch 默认 FP32，导出 ONNX 时若未显式转为 FP16，后续 TensorRT 的 FP16 优化可能因精度不匹配产生数值偏差。应在导出前统一 model.half()。
  4. 动态轴设置错误：ONNX 导出时若未指定 dynamic_axes，生成的图会固定 batch size 和序列长度，导致推理时无法处理其他形状。
  5. 跳过验证直接上线：转换后必须对比原始 PyTorch 模型与 TensorRT engine 的输出 logits，使用 torch.allclose 设置 atol=1e-2（FP16 下）验证一致性。
复盘建议：在项目启动时，先明确“最终部署硬件型号”和“TensorRT 版本”，再决定转换路径。对于 LLM，优先选择 TensorRT-LLM 原生工具链；对于 CNN 或小型模型，ONNX 中转仍然高效。每次转换后保留完整的转换日志与版本号，便于回溯问题。
← 上一章 📖 返回目录 下一章 →