📚 课程目录第4章 TensorRT-LLM 单机部署与精度验证

4. TensorRT-LLM 单机部署与精度验证

任务目标

在单台NVIDIA GPU服务器上完成TensorRT-LLM的部署，将Meta Llama 3.1 8B模型转换为TensorRT引擎，并通过对比TensorRT引擎与HuggingFace原始模型在相同输入下的输出，验证TensorRT-LLM推理结果的数值一致性。

环境与依赖

操作系统：Ubuntu 22.04+，NVIDIA驱动≥535，CUDA 12.4
GPU：单卡，显存≥24GB（推荐A100/A10G/L40S）
Docker镜像：nvcr.io/nvidia/tensorrt-llm:24.08
Python依赖：transformers、torch（镜像内已预装）
模型：meta-llama/Llama-3.1-8B-Instruct（需HuggingFace访问令牌）

实现步骤

1. 启动容器并转换模型

使用官方镜像启动容器，挂载模型缓存目录，然后执行模型转换脚本生成TensorRT引擎。转换过程将HuggingFace权重转换为FP16精度的TensorRT引擎文件。

# 启动容器
docker run --gpus all -it --shm-size=2g \
  -v /path/to/models:/models \
  -v /path/to/output:/trt_engine \
  nvcr.io/nvidia/tensorrt-llm:24.08

# 容器内执行模型转换
python /app/tensorrt_llm/examples/llama/convert_checkpoint.py \
  --model_dir /models/meta-llama/Llama-3.1-8B-Instruct \
  --output_dir /trt_engine/llama-8b-fp16 \
  --dtype float16

# 构建TensorRT引擎
trtllm-build \
  --checkpoint_dir /trt_engine/llama-8b-fp16 \
  --output_dir /trt_engine/llama-8b-engine \
  --gemm_plugin float16 \
  --max_batch_size 1 \
  --max_input_len 2048 \
  --max_seq_len 4096

2. 编写精度验证脚本

以下脚本加载同一输入，分别通过TensorRT引擎和HuggingFace模型生成输出，对比logits差异。

# verify_precision.py
import torch
import tensorrt_llm
from tensorrt_llm.runtime import ModelRunner
from transformers import AutoTokenizer, AutoModelForCausalLM

# 配置
model_path = "/models/meta-llama/Llama-3.1-8B-Instruct"
engine_dir = "/trt_engine/llama-8b-engine"
test_prompt = "The capital of France is"
max_new_tokens = 32

# 初始化tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path)
input_ids = tokenizer(test_prompt, return_tensors="pt").input_ids.cuda()

# TensorRT-LLM推理
runner = ModelRunner.from_dir(engine_dir=engine_dir, lora_dir=None)
trt_output = runner.generate(
    [test_prompt],
    max_new_tokens=max_new_tokens,
    end_id=tokenizer.eos__id,
    pad_id=tokenizer.pad__id
)
trt_tokens = trt_output[0]["output_ids"][0][len(input_ids[0]):]
trt_text = tokenizer.decode(trt_tokens, skip_special_tokens=True)

# HuggingFace推理（参考实现）
hf_model = AutoModelForCausalLM.from_pretrained(
    model_path, torch_dtype=torch.float16, device_map="cuda"
)
hf_model.eval()
with torch.no_grad():
    hf_output = hf_model.generate(
        input_ids,
        max_new_tokens=max_new_tokens,
        do_sample=False,
        pad__id=tokenizer.pad__id
    )
hf_tokens = hf_output[0][len(input_ids[0]):]
hf_text = tokenizer.decode(hf_tokens, skip_special_tokens=True)

# 对比结果
print(f"TensorRT-LLM输出: {trt_text}")
print(f"HuggingFace输出:  {hf_text}")
print(f"完全一致: {trt_text == hf_text}")

# 逐对比
trt_ids = trt_tokens.tolist()
hf_ids = hf_tokens.tolist()
mismatch = [(i, t, h) for i, (t, h) in enumerate(zip(trt_ids, hf_ids)) if t != h]
if mismatch:
    print(f"不一致数: {len(mismatch)}")
    for idx, t, h in mismatch[:5]:
        print(f"  位置{idx}: TRT={tokenizer.decode([t])} HF={tokenizer.decode([h])}")
else:
    print("所有完全一致")

运行与验证

将上述脚本保存为verify_precision.py，放入容器内工作目录。
执行python verify_precision.py。
观察输出：若打印完全一致: True且无mismatch列表，则精度验证通过。
若存在不一致，检查mismatch列表中的位置和内容——少量尾部差异（如换行符）在可接受范围内，但前几个生成必须完全一致。

判定标准

TensorRT引擎成功构建，无报错。
脚本正常运行，TensorRT-LLM与HuggingFace生成的前max_new_tokens个完全一致（或仅末尾存在因采样边界导致的微小差异）。
输出文本语义一致，无乱码或明显错误。

常见问题

显存不足：降低max_batch_size为1，或改用--dtype float16减少显存占用。
tokenizer不一致：确保TensorRT-LLM和HuggingFace使用同一tokenizer文件，且pad_设置一致。
引擎构建失败：检查CUDA版本与镜像匹配，确认模型路径正确且权重完整。