第3章模型准备

3. 模型准备：支持模型列表、模型下载、模型格式转换

模型准备这一步，说白了就是「把饭做好，等会儿才能吃」。我见过不少新手，一上来就急着跑 benchmark，结果模型格式不对，或者压根不支持，白白浪费半天时间。咱们今天就把这块捋清楚。

3.1 支持模型列表：别瞎选，先看官方清单

TensorRT-LLM 不是万能的，它只支持一部分主流模型。我个人习惯，第一步永远是去翻官方支持的模型列表。你想想看，要是选了个不支持的模型，后面所有工作都白搭。

目前官方主要支持以下几类：

LLaMA 系列：LLaMA、LLaMA 2、LLaMA 3、Code LLaMA 等
ChatGLM 系列：ChatGLM-6B、ChatGLM2、ChatGLM3
Qwen 系列：Qwen、Qwen1.5、Qwen2
Baichuan 系列：Baichuan-7B、Baichuan2
Mistral 系列：Mistral-7B、Mixtral 8x7B
Falcon 系列：Falcon-7B、Falcon-40B、Falcon-180B
GPT 系列：GPT-J、GPT-NeoX、BLOOM
其他：InternLM、Yi、DeepSeek 等

小提示：我建议你直接去 TensorRT-LLM 的 GitHub 仓库，看 examples/ 目录下的文件夹。每个模型都有对应的示例脚本，一目了然。

嗯，这里要注意：不是所有 HuggingFace 上的模型都能直接转。有些模型结构比较特殊，需要额外适配。我在项目中遇到过，有个团队想转一个冷门模型，结果折腾了两天，最后还是换成了支持的模型。

3.2 模型下载：从 HuggingFace 拉下来

确定模型支持后，下一步就是下载。我一般用 huggingface_hub 这个库，简单直接。

# 安装依赖
pip install huggingface_hub

# 下载模型（以 LLaMA-2-7B 为例）
from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="meta-llama/Llama-2-7b-chat-hf",
    local_dir="./models/Llama-2-7b-chat-hf",
    local_dir_use_symlinks=False
)

如果你在墙内，下载可能会慢。我建议用镜像站：

# 设置镜像
export HF_ENDPOINT=https://hf-mirror.com

# 或者直接在代码里指定
import os
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"

注意：有些模型（比如 LLaMA 系列）需要申请权限。你得先去 HuggingFace 上填申请表，等审核通过后才能下载。我曾经等过一周才拿到权限，所以建议提前准备。

3.3 模型格式转换：HuggingFace → TensorRT-LLM

这是最核心的一步。HuggingFace 的模型是 PyTorch 格式，TensorRT-LLM 需要的是自己特有的格式。说白了，就是把模型「翻译」成 TensorRT 能理解的语言。

转换流程大致如下：

准备权重：确保 HuggingFace 模型下载完整
运行转换脚本：TensorRT-LLM 为每个模型都提供了转换脚本
生成引擎：转换后的模型还需要编译成 TensorRT 引擎

我拿 LLaMA-2 举个例子：

# 进入 TensorRT-LLM 的 LLaMA 示例目录
cd tensorrt_llm/examples/llama

# 转换权重
python convert_checkpoint.py \
    --model_dir ./models/Llama-2-7b-chat-hf \
    --output_dir ./tllm_checkpoints/llama-7b \
    --dtype float16

# 构建 TensorRT 引擎
trtllm-build \
    --checkpoint_dir ./tllm_checkpoints/llama-7b \
    --output_dir ./tllm_engines/llama-7b \
    --gemm_plugin float16 \
    --max_batch_size 8 \
    --max_input_len 2048 \
    --max_output_len 512

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

--dtype float16：用半精度，省显存。我一般先用 float16，如果精度不够再换 float32
--gemm_plugin float16：启用矩阵乘法插件，能加速推理
--max_batch_size：最大 batch 大小，根据你的显存来设
--max_input_len / --max_output_len：输入输出最大长度，别设太大，否则显存会爆

  核心要点：转换过程中，最常遇到的问题就是显存不足。我建议你先把 max_batch_size 设小一点（比如 4），等跑通了再慢慢调大。

3.4 避坑指南：我踩过的那些坑

做模型转换，我至少踩过三次坑，分享出来大家少走弯路：

版本不匹配：TensorRT-LLM 的版本和 HuggingFace 的 transformers 版本要对应。我曾经用了一个新版的 transformers，结果转换脚本报错，最后降级才解决。
权重路径别搞错：--model_dir 要指向包含 config.json 的目录，不是模型文件夹的父目录。
显存不够怎么办：如果显存不够，可以试试 --use_weight_only 参数，用 INT8 量化，能省一半显存。

3.5 知识体系总览

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

这张图把整个流程串起来了。你从左边开始，先选模型，再下载，最后转换。每一步都有坑，但按我说的来，基本能一次过。

好了，模型准备这块就讲这么多。记住一句话：模型选对，下载完整，参数调好，后面跑 benchmark 就顺了。

我的习惯：每次转换完，我都会先跑一个简单的推理测试，确认输出正常。别等到 benchmark 跑完了才发现模型没转对，那就尴尬了。

公众号:蓝海资料掘金营，微信deep3321

📚 课程目录

3. 模型准备：支持模型列表、模型下载、模型格式转换

3.1 支持模型列表：别瞎选，先看官方清单

3.2 模型下载：从 HuggingFace 拉下来

3.3 模型格式转换：HuggingFace → TensorRT-LLM

3.4 避坑指南：我踩过的那些坑

3.5 知识体系总览