3. 模型加载与引擎构建：从HuggingFace模型到TensorRT引擎的完整流程、量化配置（FP16/INT8/INT4）

好，咱们直接进入正题。这一章，我带你走一遍从 HuggingFace 上拉下一个模型，到最终生成一个能在 GPU 上跑得飞快的 TensorRT 引擎的完整流程。说白了，就是把一个“通用”的模型，变成专为你硬件“定制”的、能榨干每一丝性能的“特供版”。

这个过程，我习惯把它拆成三步：模型准备 → 模型转换 → 引擎构建。每一步都有坑，每一步也都有优化空间。咱们一个一个来看。

3.1 第一步：模型准备——从 HuggingFace 到本地

这一步其实很简单，就是把你想要的模型从 HuggingFace 的模型库下载下来。但这里有个细节，我刚开始做的时候吃过亏。

具体怎么做？看代码：

# 1. 下载模型到本地
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "meta-llama/Llama-2-7b-chat-hf"
# 先下载，保存到本地目录
model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir="./hf_models")
tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="./hf_models")

# 2. 保存到本地，方便后续使用
model.save_pretrained("./local_llama_model")
tokenizer.save_pretrained("./local_llama_model")

嗯，这里要注意，cache_dir 参数可以指定下载路径。我个人习惯把所有模型都放在一个统一的目录下，方便管理。

3.2 第二步：模型转换——从 PyTorch 到 ONNX

TensorRT 不能直接吃 PyTorch 的模型。它需要一个中间格式——ONNX。这一步，就是把 PyTorch 模型“翻译”成 ONNX 格式。

为什么要转成 ONNX？说白了，ONNX 是一个“通用语言”，它把模型的计算图固定下来，TensorRT 才能在这个固定图上做各种优化，比如算子融合、内存复用等等。

转换过程，我直接上代码：

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("./local_llama_model")
model.eval()  # 一定要切换到评估模式

# 创建一个 dummy input，用于 trace 模型
dummy_input = torch.randint(0, 100, (1, 128))  # batch_size=1, seq_len=128

# 导出 ONNX
torch.onnx.export(
    model,
    dummy_input,
    "llama_model.onnx",
    input_names=["input_ids"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids": {0: "batch_size", 1: "seq_len"},
        "logits": {0: "batch_size", 1: "seq_len"}
    },
    opset_version=17,
    do_constant_folding=True,
)

你想想看，如果模型只支持固定长度，那用户发一个短句，你还得给它 padding 到 128，浪费了多少计算资源？所以，dynamic_axes 是必须的。

3.3 第三步：引擎构建——从 ONNX 到 TensorRT 引擎

这一步是核心中的核心。TensorRT 会根据你的 GPU 型号，把 ONNX 模型编译成一个高度优化的引擎文件（.plan 或 .engine）。这个过程，我习惯用 trtexec 命令行工具来做，因为它最直观，也最容易调试。

基本的命令长这样：

trtexec --onnx=llama_model.onnx \
        --saveEngine=llama_fp16.engine \
        --fp16 \
        --workspace=4096 \
        --minShapes=input_ids:1x1 \
        --optShapes=input_ids:1x128 \
        --maxShapes=input_ids:1x512

这里有几个参数，我解释一下：

• --fp16：启用 FP16 精度。这是最常用的量化方式，性能提升明显，精度损失几乎可以忽略。
• --workspace：指定 TensorRT 构建引擎时可用的最大显存，单位是 MB。这个值设得越大，TensorRT 就有越多的空间去做优化，但也要注意别把显存撑爆了。
• --minShapes/optShapes/maxShapes：这三个参数定义了动态 shape 的范围。TensorRT 会根据你提供的 optShapes 做针对性优化，所以 optShapes 要尽量接近你线上服务的实际输入 shape。

3.4 量化配置：FP16 / INT8 / INT4 怎么选？

量化，说白了就是用更少的 bit 来表示模型参数，从而减少显存占用和计算量。但代价是精度会下降。怎么选？我直接给你一个表格，一目了然：

量化类型	显存节省	推理加速	精度损失	适用场景
FP16	约 50%	约 1.5x - 2x	几乎无	绝大多数场景，首选
INT8	约 75%	约 2x - 3x	轻微	对精度要求不高的场景，如简单对话
INT4	约 87.5%	约 3x - 4x	明显	资源极度受限，或对精度容忍度高的场景

我个人建议，FP16 是底线。除非你的显存实在不够，或者对延迟有极致要求，否则不要轻易尝试 INT8 或 INT4。因为一旦精度损失过大，模型输出的质量会明显下降，用户是能感知到的。

INT8 和 INT4 的构建命令，只需要在 trtexec 中替换 --fp16 为 --int8 或 --int4 即可。但要注意，INT8 和 INT4 通常需要校准数据集（calibration dataset）来减少量化误差。这个校准数据集，最好是从你的真实业务数据中采样出来的。

3.5 完整流程的 SVG 框架图

为了让你更直观地理解整个流程，我画了一张图：

这张图把整个流程串起来了。从 HuggingFace 下载，到本地保存，再到 ONNX 转换，最后通过 trtexec 编译成 TensorRT 引擎。中间的量化选项，就是你根据业务需求做的取舍。

3.6 总结与个人心得

好了，这一章的内容就这么多。从模型准备到引擎构建，每一步都有它的门道。我最后再唠叨几句：

• 不要跳过 ONNX 这一步。虽然 TensorRT 现在也支持直接加载 PyTorch 模型（通过 torch-tensorrt），但 ONNX 的兼容性和调试便利性是最好的。我遇到问题，第一反应就是导出 ONNX，然后用 onnxruntime 跑一遍，看看是不是模型本身的问题。
• 量化要谨慎。FP16 是安全牌，INT8 和 INT4 是“富贵险中求”。如果你没有足够的业务数据来做校准，或者对精度要求很高，那就老老实实用 FP16。
• 引擎构建是“一次性”的。构建过程可能很慢（大模型可能要几个小时），但一旦构建完成，引擎文件就可以反复使用。所以，构建时多花点时间优化参数，是值得的。

嗯，这一章就到这里。下一章，我们会聊聊如何用 TensorRT-LLM 的 Python API 来加载这个引擎，并跑通第一个推理请求。到时候见。

---