3、模型转换基础：理解ONNX、TorchScript、FasterTransformer等模型格式，以及转换为TensorRT-LLM的流程

3.1 为什么需要关注模型格式

在进入TensorRT-LLM编译之前，首先要回答一个关键问题：我们手里拿到的模型文件，到底是以什么格式存在的？不同训练框架产出的模型格式，决定了后续编译管线的入口。PyTorch训练出的模型通常是.pt或.pth（包含权重与网络结构），HuggingFace生态则统一使用safetensors + config.json的目录结构。而TensorRT-LLM本身并不直接消费这些原始格式，它需要的是经过特定转换后的中间表示。理解这一点，是避免“拿到模型却无从下手”的第一步。

核心认知：TensorRT-LLM的编译入口不是PyTorch权重文件，而是转换后的模型格式。整个转换流程的本质，是将训练框架的模型描述“翻译”为TensorRT-LLM能够识别的引擎构建输入。

3.2 三种主流模型格式的定位与差异

在TensorRT-LLM的生态中，你会频繁接触到ONNX、TorchScript和FasterTransformer（现已并入TensorRT-LLM）三种格式。它们各自服务于不同的转换阶段，不能混为一谈。

格式 本质 在TensorRT-LLM中的角色 典型来源
ONNX 跨框架的静态计算图 作为中间桥梁，将PyTorch模型导出为通用图，再转换为TensorRT-LLM的权重格式 torch.onnx.export()
TorchScript PyTorch的序列化图表示 较少直接用于TensorRT-LLM，更多用于PyTorch自身推理优化；在转换流程中可作为备选导出路径 torch.jit.trace/script
FasterTransformer 早期高性能推理的权重/引擎格式 其权重格式与TensorRT-LLM兼容，部分旧模型可直接复用其转换脚本 FasterTransformer仓库的utils脚本

需要特别强调的是：ONNX不是TensorRT-LLM的最终引擎格式。ONNX在这里扮演的是“通用交换格式”的角色。TensorRT-LLM真正需要的是权重文件（通常是.safetensors或.bin）加模型配置，然后通过其内部的trtllm-build命令直接构建引擎。因此，ONNX导出往往只是转换管线中的一个中间步骤，而非终点。

3.3 从PyTorch到TensorRT-LLM的标准转换流程

下面以最常见的HuggingFace模型为例，梳理一条完整的转换路径。整个流程分为四个阶段，每个阶段都有明确的输入输出。

# 阶段1：导出ONNX（可选但推荐用于验证）
import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
dummy_input = torch.randint(0, 32000, (1, 64))
torch.onnx.export(
    model, dummy_input, "llama2-7b.onnx",
    opset_version=17,
    input_names=["input_ids"],
    output_names=["logits"]
)

# 阶段2：使用官方转换脚本生成TensorRT-LLM权重
python convert_checkpoint.py \
    --model_dir ./llama2-7b-hf \
    --output_dir ./tllm_checkpoint \
    --dtype float16

# 阶段3：构建TensorRT引擎
trtllm-build \
    --checkpoint_dir ./tllm_checkpoint \
    --output_dir ./llama2-7b-engine \
    --gemm_plugin float16

# 阶段4：加载引擎进行推理
from tensorrt_llm.runtime import ModelRunner
runner = ModelRunner.from_dir(engine_dir="./llama2-7b-engine")

上述流程中，阶段2的convert_checkpoint.py是核心。它读取HuggingFace的config.json和safetensors权重，将其重新组织为TensorRT-LLM自定义的checkpoint格式。这个格式包含了模型结构参数（层数、头数、维度等）以及经过重排的权重张量。值得注意的是，ONNX导出在官方推荐流程中并非必需——TensorRT-LLM的转换脚本直接解析HuggingFace格式，跳过了ONNX这一步。ONNX更多用于调试、可视化或对接其他推理后端。

实操建议：如果你的模型在HuggingFace上有官方支持（如Llama、Qwen、ChatGLM等），直接使用TensorRT-LLM仓库中对应的convert_checkpoint.py即可，无需手动导出ONNX。只有当模型结构较新、官方脚本尚未支持时，才考虑通过ONNX作为中间表示进行自定义转换。

3.4 案例：Llama-2-7B的完整转换复盘

我们以Llama-2-7B为例，走一遍真实转换过程，并记录关键输出。

# 1. 准备HuggingFace模型（已下载至本地）
ls ./llama2-7b-hf/
# config.json  generation_config.json  pytorch_model.bin  tokenizer.json  tokenizer_config.json

# 2. 执行转换（关键日志）
python convert_checkpoint.py --model_dir ./llama2-7b-hf --output_dir ./tllm_checkpoint --dtype float16
# [INFO] Loading model from ./llama2-7b-hf
# [INFO] Model config: LlamaConfig(vocab_size=32000, hidden_size=4096, num_hidden_layers=32, num_attention_heads=32)
# [INFO] Converting 32 decoder layers...
# [INFO] Saving checkpoint to ./tllm_checkpoint
# [INFO] Done. Total 128 files written.

# 3. 构建引擎
trtllm-build --checkpoint_dir ./tllm_checkpoint --output_dir ./llama2-7b-engine --gemm_plugin float16
# [INFO] Engine built in 45.3 seconds.
# [INFO] Engine size: 13.4 GB (FP16)

这个案例揭示了几个关键点：

  • 转换时间远小于构建时间：checkpoint转换只是权重重排，而引擎构建涉及算子选择、内核调优和显存规划。
  • 权重文件数量众多：128个文件对应每一层的QKV投影、输出投影、MLP等独立张量，TensorRT-LLM按层分文件存储，便于并行加载。
  • FP16引擎体积约为原始FP32权重的1/2：因为权重本身转为FP16存储，且引擎中包含了kernel代码和调度元数据。

3.5 常见误区与排障复盘

在模型转换环节，初学者最容易踩入以下几个坑。我们逐一复盘并给出对策。

误区一：认为ONNX是必经之路。实际上，TensorRT-LLM官方转换脚本直接读取HuggingFace格式，ONNX导出反而可能引入额外的算子兼容性问题。除非你要做跨框架迁移或可视化调试，否则不必导出ONNX。
误区二：忽略dtype参数的影响。转换时指定的--dtype float16决定了权重存储精度。如果模型有某些层对精度敏感（如LayerNorm的beta/gamma），可能需要使用float32保留这些层，否则推理精度会下降。建议转换前检查模型配置中的torch_dtype字段。
误区三：混淆“转换”与“构建”。很多同学把convert_checkpoint.py和trtllm-build混为一谈。实际上，前者只是把HuggingFace权重“翻译”成TensorRT-LLM的checkpoint格式，不涉及任何GPU优化；后者才真正生成针对特定GPU的优化引擎。如果换了GPU型号，只需要重新执行trtllm-build，无需重新转换权重。

另一个常见问题是版本不匹配。TensorRT-LLM对模型结构的支持是硬编码的，如果你的模型是Llama-3.2而本地TensorRT-LLM版本只支持到Llama-3.1，转换脚本会直接报错或生成错误权重。此时需要：

  1. 检查tensorrt_llm/models/目录下是否有对应模型定义文件；
  2. 若无，考虑升级TensorRT-LLM版本，或使用--model_type参数指定相近的模型结构；
  3. 最稳妥的方案是查阅官方文档的“Supported Models”列表，确认模型版本在支持范围内。
TensorRT-LLM 模型转换全景流程 HuggingFace 模型 config.json + *.safetensors (或 pytorch_model.bin) TensorRT-LLM Checkpoint 每层独立权重文件 (约128个 .bin 文件) TensorRT 引擎 .engine 文件 (含kernel + 调度元数据) convert_checkpoint.py --dtype float16 trtllm-build --gemm_plugin float16 ONNX 中间表示 (可选调试/验证用) torch.onnx.export() 自定义解析脚本 运行时推理 ModelRunner.from_dir() ⚠ 常见误区标注 ① ONNX不是必经之路——官方转换脚本直接读取HuggingFace格式 ② 转换(convert) ≠ 构建(build)——前者仅重排权重，后者才生成GPU优化引擎 ③ 换GPU型号只需重新 trtllm-build，无需重新 convert_checkpoint ④ 版本不匹配时优先查 Supported Models 列表，而非盲目升级 完整链路：HuggingFace → TensorRT-LLM Checkpoint → TensorRT Engine → 运行时推理

3.6 转换流程中的关键检查清单

为了确保转换过程顺利，建议在动手前逐项核对以下清单：

  • 模型版本确认：在TensorRT-LLM官方文档的Supported Models列表中确认你的模型架构和版本受支持。
  • 权重精度规划：明确目标精度（FP16/BF16/FP8），并检查模型是否有特殊层需要保留FP32。
  • 分词器完整性：转换脚本通常不处理tokenizer，但推理时需要。确保tokenizer.json或tokenizer.model与模型权重放在同一目录。
  • 磁盘空间：转换后的checkpoint体积约为原始权重的1~1.5倍（取决于dtype），引擎文件另占空间。预留至少2倍模型大小的磁盘。
  • GPU显存预检：构建引擎时会将checkpoint加载到显存，建议显存大于模型FP16体积的1.5倍。

完成上述检查后，转换流程通常可以一次通过。如果遇到报错，优先查看convert_checkpoint.py输出的日志尾部——绝大多数错误会明确提示是“模型结构不支持”还是“权重文件缺失”，据此对症下药即可。

← 上一章 📖 返回目录 下一章 →