模型导出与转换：从HuggingFace到TensorRT-LLM

模型导出与转换，说白了就是让模型换个格式、换个环境跑起来。我刚开始接触TensorRT-LLM时，觉得这步不就是个格式转换嘛，有啥难的？结果第一次就栽了跟头——导出的模型跑起来全是乱码。嗯，这里面的门道还真不少。

今天咱们就聊聊，怎么把HuggingFace上那些大模型，一步步变成TensorRT-LLM能用的格式。我个人习惯把这过程拆成四步：导出ONNX、转成Checkpoint、处理权重绑定、最后分片。每一步都有坑，咱们一个一个填。

从HuggingFace导出模型

HuggingFace上的模型，通常以PyTorch的bin文件或safetensors格式存储。我建议直接用transformers库加载，别自己手写解析——我之前试过，结果被各种命名规则搞晕了。

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "meta-llama/Llama-2-7b-hf"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 保存为safetensors格式（推荐）
model.save_pretrained("./llama2-7b-hf", safe_serialization=True)
tokenizer.save_pretrained("./llama2-7b-hf")

ONNX导出：中间表示的重要性

ONNX说白了就是个中间格式。它不直接用于推理，但能帮我们做两件事：一是检查模型结构对不对，二是做图优化。我个人习惯在转TensorRT-LLM之前，先过一遍ONNX。

为什么要先导出ONNX？你想想看，HuggingFace的模型里有很多动态控制流（比如if-else），这些在TensorRT里跑不了。ONNX导出时会把它们展开成静态图，提前暴露问题。

import torch
import onnx
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("./llama2-7b-hf")
dummy_input = torch.randint(0, 32000, (1, 128))

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

# 验证ONNX模型
onnx_model = onnx.load("llama2-7b.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX导出成功！")

TensorRT-LLM Checkpoint格式

TensorRT-LLM有自己的Checkpoint格式，跟HuggingFace的不一样。它把权重按层、按张量并行度重新组织，方便后续构建推理引擎。

转换工具在tensorrt_llm/examples/llama/convert_checkpoint.py里。我建议直接用这个脚本，别自己写——里面处理了很多细节，比如RoPE的权重转换、LayerNorm的合并等。

# 转换命令示例
python convert_checkpoint.py \
    --model_dir ./llama2-7b-hf \
    --output_dir ./llama2-7b-trllm \
    --dtype float16 \
    --tp_size 2 \
    --pp_size 1

参数	说明	我的建议
--model_dir	HuggingFace模型路径	确保包含config.json
--output_dir	输出Checkpoint路径	用有意义的命名
--dtype	权重精度：float16/bfloat16/float32	一般用float16
--tp_size	张量并行度	等于GPU数量
--pp_size	流水线并行度	单机一般设为1

权重绑定：把散落的碎片拼起来

权重绑定是TensorRT-LLM里一个关键操作。HuggingFace的模型里，Q、K、V的权重是分开存的，Gate和Up也是分开的。但TensorRT-LLM为了推理效率，会把它们合并成一个矩阵。

举个例子，LLaMA的Attention层：

• HuggingFace：q_proj.weight、k_proj.weight、v_proj.weight 三个独立的权重
• TensorRT-LLM：合并成 qkv.weight 一个矩阵

这样做的好处是减少内存访问次数。GPU从显存读数据是很慢的，一次读一个大矩阵比分三次读小矩阵快得多。

# 权重绑定伪代码
import torch

# 假设从HuggingFace加载的权重
q_weight = torch.load("q_proj.weight")
k_weight = torch.load("k_proj.weight")
v_weight = torch.load("v_proj.weight")

# 合并为单个矩阵
qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0)

# 保存为TensorRT-LLM格式
torch.save(qkv_weight, "qkv.weight")

权重分片：让多卡并行跑起来

权重分片是为了张量并行。你想想看，如果模型有70B参数，一张A100放不下，那就得切到多张卡上。TensorRT-LLM支持按行切分和按列切分两种方式。

分片的规则其实很简单：

• 按行切分：把权重矩阵的行切给不同GPU，适用于Linear层的输入变换
• 按列切分：把权重矩阵的列切给不同GPU，适用于Linear层的输出变换

转换脚本会自动处理这些。你只需要指定--tp_size，它就会把每个权重切成tp_size份。

# 分片后的目录结构
llama2-7b-trllm/
├── tp_rank_0/
│   ├── model.layers.0.attention.qkv.weight
│   ├── model.layers.0.mlp.gate_up.weight
│   └── ...
├── tp_rank_1/
│   ├── model.layers.0.attention.qkv.weight
│   ├── model.layers.0.mlp.gate_up.weight
│   └── ...
└── config.json

总结一下

模型导出与转换，说白了就是让模型换个活法。从HuggingFace到ONNX，再到TensorRT-LLM Checkpoint，每一步都有它的意义。我个人觉得，最关键的还是理解权重绑定和分片——这两个操作直接影响推理性能。

嗯，今天就聊到这儿。记住一点：转换过程中多验证，别等到跑推理才发现问题。我吃过这个亏，希望你不用再吃一遍。