第3章模型仓库与转换

一、模型仓库与转换：HuggingFace 模型下载、ONNX 导出、TensorRT-LLM 检查点格式

好，咱们正式开始。这一章，我带你走一遍模型编译前的“原料准备”流程。说白了，就是搞清楚三件事：模型从哪来、怎么转成中间格式、最终喂给 TensorRT-LLM 的是什么样子。

我个人习惯把这一步叫做“数据管道的前置工程”。很多同学一上来就急着编译，结果卡在模型格式不兼容上，一查就是半天。嗯，咱们先把地基打牢。

1.1 HuggingFace 模型下载：你的模型仓库入口

HuggingFace 现在是事实上的模型集散地。你想想看，从 LLaMA、Mistral 到 Qwen，几乎所有的开源大模型都能在这找到。下载模型其实就一行命令的事，但有几个坑我得提前跟你说。

核心命令：

pip install huggingface-hub
huggingface-cli download meta-llama/Llama-2-7b-chat-hf --local-dir ./llama2-7b

这里要注意几点：

需要登录：像 LLaMA 这种模型有访问限制。你得先去 HuggingFace 官网申请权限，然后用 huggingface-cli login 登录。我曾经因为忘了这步，脚本跑了一小时报错，气得不行。
断点续传：模型动不动几十 G，网络断了很头疼。加个 --resume-download 参数，断了能接着下。
只下载必要文件：如果你只需要模型权重，用 --include "*.safetensors" 过滤，别把一堆 tokenizer 配置也拖下来。

我的小技巧： 下载前先看看模型卡页面的 Files and versions 标签。有些模型有多个分支（比如 main 和 fp16），选对版本能省不少时间。

1.2 ONNX 导出：打通框架的桥梁

ONNX 是什么？说白了，它就是模型界的“普通话”。不管你是 PyTorch、TensorFlow 还是 JAX 训练的模型，都能转成 ONNX 格式。TensorRT-LLM 虽然原生支持很多格式，但 ONNX 依然是最通用的中间桥梁。

导出 ONNX 的典型做法是这样的：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("./llama2-7b")
tokenizer = AutoTokenizer.from_pretrained("./llama2-7b")

# 构造一个 dummy input
dummy_input = tokenizer("Hello, how are you?", return_tensors="pt")

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

这里有个关键点：dynamic_axes。你想想看，推理时 batch size 和序列长度都是变化的，如果不设置动态轴，导出的 ONNX 就只能处理固定尺寸的输入。我在项目中遇到过，有人忘了设这个，结果推理时输入稍微长一点就报错。

注意： ONNX 导出时，opset_version 不是越高越好。TensorRT-LLM 对某些高版本算子支持可能不完善。我个人习惯用 opset 17，兼容性最好。

1.3 TensorRT-LLM 检查点格式：编译前的最后一步

好，现在模型已经转成 ONNX 了。但 TensorRT-LLM 并不直接吃 ONNX，它需要一种叫 检查点（Checkpoint） 的中间格式。你可以把它理解成“TensorRT-LLM 能读懂的营养餐”。

为什么要多这一步？因为 TensorRT-LLM 需要对模型做很多优化，比如层融合、量化、内存重排。这些优化需要知道模型的详细结构信息，而 ONNX 文件里这些信息是“散装”的。检查点格式把这些信息整理成 TensorRT-LLM 方便处理的格式。

转换命令很简单：

# 假设你已经有了 HuggingFace 格式的模型
python examples/llama/convert_checkpoint.py \
    --model_dir ./llama2-7b \
    --output_dir ./llama2-7b-trt-ckpt \
    --dtype float16

这个脚本会做几件事：

读取 HuggingFace 的 config.json 和权重文件
重新组织权重布局（比如把 QKV 矩阵合并）
按 TensorRT-LLM 的格式保存为 .bin 文件

检查点目录结构示例：

llama2-7b-trt-ckpt/
├── config.json
├── rank_0/
│   ├── model.layers.0.self_attn.qkv.weight.bin
│   ├── model.layers.0.self_attn.o_proj.weight.bin
│   └── ...
└── tokenizer/

看到 rank_0 这个目录没？这说明 TensorRT-LLM 天然支持多卡并行。如果你有 4 张卡，转换时会生成 rank_0 到 rank_3 四个目录，每个目录里放的是分片后的权重。

避坑指南： 我曾经遇到一个情况，转换时指定了 --dtype float16，但模型里有些层是 int8 的，结果转换报错。后来发现是 HuggingFace 模型里混用了精度。解决办法是先用 torch.save 把所有权重统一成 float16 再转。

1.4 知识体系总览

为了让你对整个流程有个直观印象，我画了张图：

这张图把整个流程串起来了。你可能会问：为什么 ONNX 那一步画成虚线？因为在实际工作中，有些模型可以直接从 HuggingFace 格式跳到 TensorRT-LLM 检查点，不需要经过 ONNX。比如 LLaMA 系列，官方脚本已经支持直接转换。但 ONNX 作为通用格式，理解它对你调试自定义模型很有帮助。

1.5 本章小结

好，这一章的内容就到这。总结一下：

HuggingFace 下载：注意登录、断点续传、文件过滤
ONNX 导出：动态轴是关键，opset 版本别太高
TRT-LLM 检查点：理解目录结构，注意精度统一

下一章，我们就要真正开始编译了。到时候你会看到，前面这些准备工作做得越扎实，编译过程就越顺畅。嗯，今天就到这，有问题随时问我。

📚 课程目录

一、模型仓库与转换：HuggingFace 模型下载、ONNX 导出、TensorRT-LLM 检查点格式

1.1 HuggingFace 模型下载：你的模型仓库入口

1.2 ONNX 导出：打通框架的桥梁

1.3 TensorRT-LLM 检查点格式：编译前的最后一步

1.4 知识体系总览

1.5 本章小结