Chapter 10

TensorRT-LLM 集成:训完直接出 engine

📌 commit 2ea3e0fNeMo → TRT-LLM → Triton 部署链路

NVIDIA 训练栈的最后一步是把 NeMo .nemo 模型转成 TensorRT-LLM engine,再喂给 Triton 推理服务器。本章把这条链路串起来。

10.1整体流程

NeMo .nemo  ──▶  HF checkpoint  ──▶  TRT-LLM engine  ──▶  Triton server
   (训练)         (中间格式)          (优化的 .engine)     (HTTP API)

10.2nemo → HF

看 ch12 详细讲,这里给一个简版:

python scripts/checkpoint_converters/convert_llama_nemo_to_hf.py \
    --input_name_or_path ./llama3_8b.nemo \
    --output_path ./llama3_8b_hf

10.3HF → TRT-LLM

# 在 TensorRT-LLM 仓库内
python examples/llama/convert_checkpoint.py \
    --model_dir ./llama3_8b_hf \
    --output_dir ./trt_ckpt \
    --dtype bfloat16 \
    --tp_size 2

# 编译 engine
trtllm-build \
    --checkpoint_dir ./trt_ckpt \
    --output_dir ./trt_engine \
    --gemm_plugin bfloat16 \
    --max_input_len 4096 --max_output_len 1024 \
    --max_batch_size 32 \
    --paged_kv_cache enable \
    --use_paged_context_fmha enable

编译一次约 5-10 分钟,产出 ./trt_engine/*.engine

10.4量化选项

方案位宽典型加速命令
bf16161.0× baseline默认
FP881.8-2.2×--quant_ckpt_path ... --use_fp8(H100/H200)
INT8 SmoothQuant81.5-1.8×--smoothquant 0.5
INT4 AWQ42.5-3.5×先 AWQ 量化,再 build
INT4 GPTQ42.5-3.5×先 GPTQ 量化,再 build

10.5跑起 Triton

tritonserver --model-repository=./model_repository

# model_repository 结构
model_repository/
├── llama3-8b/
│   ├── 1/
│   │   └── trt_engine/*.engine
│   └── config.pbtxt

测试:

curl -X POST http://localhost:8000/v2/models/llama3-8b/generate \
    -d '{"text_input":"你好","max_tokens":50}'

10.6性能对比

Llama-3-8B、A100-80GB、batch=32、prompt=512、output=128:

方案tokens/slatency p50
HF transformers (bf16)~3,500~85 ms
vLLM (bf16)~12,000~32 ms
TRT-LLM (bf16)~14,000~28 ms
TRT-LLM (FP8, H100)~26,000~16 ms
TRT-LLM (INT4 AWQ)~32,000~12 ms

10.7NIM:NVIDIA 一体化部署

2024+ NVIDIA 推出 NIM(NVIDIA Inference Microservices),把 TRT-LLM + Triton + 优化全包,开箱即用 Docker:

docker run --gpus all -p 8000:8000 nvcr.io/nim/meta/llama-3-8b-instruct:1.0.0

对自家 finetune 模型不适用,但对开源公开模型是最快上线方式。

10.8常见错误

现象处理
convert_checkpoint 报 shape mismatchtp_size 不对;要和后续 build 一致
build 时 OOM降 max_batch_size;分 build_config 优化器
跑起来精度变差FP8 / INT8 都有精度损失;先 bf16 验证,再量化
Triton 启动 model 加载失败config.pbtxt 里的 backend / parameters 没对
多卡 TP 不行TRT-LLM 要 NCCL;查 NCCL 环境

10.9这章你需要带走的