Chapter 10
TensorRT-LLM 集成:训完直接出 engine
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量化选项
| 方案 | 位宽 | 典型加速 | 命令 |
|---|---|---|---|
| bf16 | 16 | 1.0× baseline | 默认 |
| FP8 | 8 | 1.8-2.2× | --quant_ckpt_path ... --use_fp8(H100/H200) |
| INT8 SmoothQuant | 8 | 1.5-1.8× | --smoothquant 0.5 |
| INT4 AWQ | 4 | 2.5-3.5× | 先 AWQ 量化,再 build |
| INT4 GPTQ | 4 | 2.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/s | latency 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 mismatch | tp_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这章你需要带走的
- NeMo → HF → TRT-LLM → Triton 是 NVIDIA 一站式链路;
- TRT-LLM 比 vLLM 略快、比 transformers 快 4-8×;
- FP8(H100)和 INT4 AWQ 是主要量化方式,能再 2× 加速;
- 开箱即用的公开模型部署用 NIM Docker,自家 finetune 走完整链路。