2. 模型转换基础

任务目标

将HuggingFace上的小型模型（）转换为TensorRT-LLM支持的格式，并完成TensorRT引擎构建，最终验证转换后的模型可以正常加载并执行推理。

环境与依赖

本案例基于TensorRT-LLM官方Docker镜像（nvcr.io/nvidia/tritonserver:24.05-trtllm-python-py3）运行，镜像内已包含TensorRT-LLM及所需依赖。宿主机需安装NVIDIA驱动（≥535）并支持CUDA 12.x。若在自定义环境运行，需安装tensorrt-llm、transformers、torch及sentencepiece。

实现步骤

整体流程分为三步：下载HuggingFace模型并转换为TensorRT-LLM的checkpoint格式，然后使用trtllm-build命令构建引擎，最后通过Python API加载引擎验证输出。

关键完整代码

以下脚本完成从HuggingFace模型到TensorRT-LLM checkpoint的转换，并自动执行引擎构建。将脚本保存为convert_and_build.py：

#!/usr/bin/env python3
import subprocess
import sys
import os
from pathlib import Path

# 1. 配置路径
hf_model_dir = Path("./hf_")          # HuggingFace模型缓存目录
trt_ckpt_dir = Path("./trt__ckpt")    # TensorRT-LLM checkpoint输出目录
engine_dir = Path("./trt__engine")    # 引擎输出目录
model_name = ""                        # 小型模型（124M参数）

# 2. 下载HuggingFace模型（若不存在）
if not hf_model_dir.exists():
    print(f"[1/3] 下载HuggingFace模型: {model_name}")
    from transformers import GPT2LMHeadModel, GPT2Tokenizer
    model = GPT2LMHeadModel.from_pretrained(model_name)
    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
    model.save_pretrained(hf_model_dir)
    tokenizer.save_pretrained(hf_model_dir)
    print(f"模型已保存至 {hf_model_dir}")

# 3. 转换为TensorRT-LLM checkpoint格式
print("[2/3] 转换为TensorRT-LLM checkpoint...")
convert_cmd = [
    sys.executable, "-m", "tensorrt_llm.models.modeling_utils",
    "--model_type", "",
    "--model_dir", str(hf_model_dir),
    "--output_dir", str(trt_ckpt_dir),
    "--dtype", "float16",
    "--tp_size", "1",
    "--pp_size", "1"
]
subprocess.run(convert_cmd, check=True)
print(f"Checkpoint已生成至 {trt_ckpt_dir}")

# 4. 构建TensorRT引擎
print("[3/3] 构建TensorRT引擎...")
build_cmd = [
    "trtllm-build",
    "--checkpoint_dir", str(trt_ckpt_dir),
    "--output_dir", str(engine_dir),
    "--gemm_plugin", "float16",
    "--max_batch_size", "1",
    "--max_input_len", "128",
    "--max_seq_len", "256"
]
subprocess.run(build_cmd, check=True)
print(f"引擎已生成至 {engine_dir}")
print("转换与构建完成。")

运行方式

在TensorRT-LLM Docker容器内执行以下命令：

python convert_and_build.py

脚本将自动完成模型下载、格式转换和引擎构建。执行时间取决于GPU性能，通常约2-5分钟。

验证方法

构建完成后，使用以下Python代码加载引擎并执行一次推理，验证模型输出是否合理：

# 保存为 verify_engine.py
import tensorrt_llm
from tensorrt_llm.runtime import ModelRunnerCpp
from transformers import GPT2Tokenizer

# 加载tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("./hf_")
tokenizer.pad_ = tokenizer.eos_

# 初始化runner
runner = ModelRunnerCpp.from_dir(
    engine_dir="./trt__engine",
    lora_dir=None,
    rank=0
)

# 构造输入
input_text = "The capital of France is"
input_ids = tokenizer.encode(input_text, return_tensors="pt")
input_lengths = [input_ids.shape[1]]

# 执行推理
outputs = runner.generate(
    batch_input_ids=[input_ids[0].tolist()],
    max_new_tokens=10,
    end_id=tokenizer.eos__id,
    pad_id=tokenizer.pad__id
)

# 解码输出
output_ids = outputs[0]["output_ids"].tolist()
generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(f"输入: {input_text}")
print(f"生成: {generated_text}")

运行验证脚本：

python verify_engine.py

通过条件：程序无报错退出，且输出文本为通顺的英文续写（例如输入“The capital of France is”时，输出应包含“Paris”或语义合理的城市名）。若输出为乱码或推理报错，则说明转换或构建环节存在问题。

常见问题

• 显存不足：若构建引擎时出现OOM，可减小--max_batch_size至1，或降低--max_seq_len至128。
• 版本不匹配：确保transformers库版本≥4.30，否则权重加载可能失败。
• 路径权限：容器内运行需确保输出目录有写入权限，建议在/workspace下执行。