TensorRT-LLM 大模型推理加速实战指南
1. 简介
TensorRT-LLM 是由 NVIDIA 发布的高性能推理库,专为大型语言模型(LLM)设计。它提供了易于使用的 Python API,允许用户定义 LLM 并构建包含最先进优化的 TensorRT 引擎,从而在 NVIDIA GPU 上实现高效的推理执行。相比通用推理框架,TensorRT-LLM 针对 Transformer 架构进行了深度优化,支持量化、连续批处理等高级特性。
TensorRT-LLM 是 NVIDIA 推出的高性能大模型推理库,本文详细介绍了其安装、引擎构建及基于 Triton 的部署流程。内容涵盖环境配置、Qwen 模型适配、Inflight Batching 优化及常见网络问题排查,旨在帮助开发者实现低延迟、高吞吐的 LLM 推理服务。

TensorRT-LLM 是由 NVIDIA 发布的高性能推理库,专为大型语言模型(LLM)设计。它提供了易于使用的 Python API,允许用户定义 LLM 并构建包含最先进优化的 TensorRT 引擎,从而在 NVIDIA GPU 上实现高效的推理执行。相比通用推理框架,TensorRT-LLM 针对 Transformer 架构进行了深度优化,支持量化、连续批处理等高级特性。
建议直接使用 Docker 环境进行部署,以避免本地环境配置冲突。首先需要安装 NVIDIA TensorRT。
sudo dpkg -i nv-tensorrt-local-repo-ubuntu2204-8.6.1-cuda-12.0_1.0-1_amd64.deb
sudo cp /var/nv-tensorrt-local-repo-ubuntu2204-8.6.1-cuda-12.0/nv-tensorrt-local-42B2FC56-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get install tensorrt
通过 pip 安装预发布版本:
pip3 install tensorrt_llm -U --pre --extra-index-url https://pypi.nvidia.com
python3 -c "import tensorrt_llm"
注意:TensorRT-LLM 与 TGI(Text Generation Inference)软件版本可能存在冲突,请确保环境隔离。
为了在生产环境中部署,通常结合 Triton Inference Server 使用。
docker pull nvcr.io/nvidia/tritonserver:24.01-trtllm-python-py3
要为现有模型创建 TensorRT 引擎,主要包含三个步骤:下载预训练权重、构建完全优化的模型引擎、部署引擎。
由于目前某些特定模型(如 Qwen2)可能尚未完全支持,示例中使用 Qwen-7B 模型。如果模型不在仓库 examples 目录下支持列表中,需要自行编写代码实现模型转化功能。
python build.py --hf_model_dir /home/chuan/models/qwen/Qwen-7B-Chat-Int4 \
--quant_ckpt_path /home/chuan/models/qwen/Qwen-7B-Chat-Int4 \
--dtype float16 \
--remove_input_padding \
--use_gpt_attention_plugin float16 \
--enable_context_fmha \
--use_gemm_plugin float16 \
--use_weight_only \
--weight_only_precision int4_gptq \
--per_group \
--world_size 1 \
--tp_size 1 \
--output_dir /home/chuan/models/qwen/Qwen-7B-Chat-Int4/trt_engines/int4-gptq/1-gpu
参数说明:
--hf_model_dir: HuggingFace 模型目录路径。--quant_ckpt_path: 量化检查点路径。--dtype: 数据类型,通常为 float16。--use_weight_only: 启用权重量化。--weight_only_precision: 量化精度,如 int4_gptq。--tp_size: 张量并行大小。--output_dir: 生成的引擎输出目录。构建完成后,可直接运行测试:
python3 ../run.py --input_text "你好,请问你叫什么?" \
--max_output_len=50 \
--tokenizer_dir /home/chuan/models/qwen/Qwen-7B-Chat-Int4 \
--engine_dir=/home/chuan/models/qwen/Qwen-7B-Chat-Int4/trt_engines/int4-gptq/1-gpu
在 tensorrtllm_backend 仓库中,all_models/inflight_batcher_llm/ 目录保存了模型执行过程不同部分的工件。该流程通常包含四个子文件夹:
复制上一步生成的模型到指定目录:
cp /home/chuan/models/qwen/Qwen-7B-Chat-Int4/trt_engines/int4-gptq/1-gpu/* /home/chuan/github/tensorrtllm_backend/all_models/inflight_batcher_llm/tensorrt_llm/1
修改 config.pbtxt,关键参数包括:
gpt_model_type: 必选,inflight_fused_batching 支持连续批处理。batch_scheduler_policy: 请求批处理策略。decoupled: 可选,是否开启流式传输。使用脚本填充模板:
python3 tools/fill_template.py --in_place \
all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt \
decoupled_mode:true,engine_dir:/all_models/inflight_batcher_llm/tensorrt_llm/1,\
gpt_model_type:inflight_fused_batching
这两个模块配置类似,主要指定 tokenizer 类型和路径。
python tools/fill_template.py --in_place \
all_models/inflight_batcher_llm/preprocessing/config.pbtxt \
tokenizer_type:auto,tokenizer_dir:/all_models/inflight_batcher_llm/tensorrt_llm/1
重新编译模型并启动 Triton Server:
python3 build.py --hf_model_dir /home/chuan/models/qwen/Qwen-7B-Chat-Int4 \
--quant_ckpt_path /home/chuan/models/qwen/Qwen-7B-Chat-Int4 \
--dtype float16 \
--remove_input_padding \
--use_gpt_attention_plugin float16 \
--enable_context_fmha \
--use_gemm_plugin float16 \
--use_weight_only \
--weight_only_precision int4_gptq \
--per_group \
--world_size 1 \
--tp_size 1 \
--output_dir /home/chuan/models/qwen/Qwen-7B-Chat-Int4/trt_engines/int4-gptq/1-gpu \
--use_inflight_batching \
--paged_kv_cache
启动 Docker 容器:
docker run -it --gpus all --network host --shm-size=2g \
-v /home/chuan/github/tensorrtllm_backend:/opt/tensorrtllm_backend \
-v /home/chuan/github/TensorRT-LLM:/opt/TensorRT-LLM \
-p 8000:8000 \
triton_trt_llm
安装依赖:
pip3 config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
pip3 install sentencepiece protobuf
启动服务器:
python3 /opt/tensorrtllm_backend/scripts/launch_triton_server.py --model_repo /opt/tensorrtllm_backend/all_models/inflight_batcher_llm --world_size 1
使用 curl 发送 POST 请求进行测试:
curl -X POST localhost:8000/v2/models/generate \
-H "Content-Type: application/json" \
-d '{"text_input": "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n你好,你叫什么?<|im_end|>\n<|im_start|>assistant\n", "max_tokens": 50}'
(注:实际请求格式需根据具体接口定义调整)
Qwen 模型可能需要单独安装的包,例如 tiktoken:
pip install tiktoken
此外,Qwen 模型对 Batch Size 有特定要求,需在 tensorrt_llm/config.pbtxt 中制定:
max_batch_size: 2
构建模型时,必须添加以下两个参数以支持连续批处理和分页 KV 缓存:
--use_inflight_batching
--paged_kv_cache
如果原本代码是为 Llama Token 设计的,而当前模型是 Qwen,则需要修改 preprocessing/model.py 和 postprocessing/model.py 中的逻辑以适配不同的 tokenizer 行为。
如果遇到 Docker 拉取镜像失败,可配置 Docker 代理。
编辑 ~/.docker/config.json:
{
"proxies": {
"default": {
"httpProxy": "http://172.17.0.1:3128",
"httpsProxy": "https://172.17.0.1:3129",
"noProxy": "*.test.example.com,.example.org,127.0.0.0/8"
}
}
}
或者通过 systemd 配置:
sudo mkdir -p /etc/systemd/system/docker.service.d
sudo vi /etc/systemd/system/docker.service.d/http-proxy.conf
内容如下:
[Service]
Environment="HTTP_PROXY=http://172.17.0.1:3128"
Environment="HTTPS_PROXY=https://172.17.0.1:3129"
Environment="NO_PROXY=localhost,127.0.0.1,docker-registry.example.com,.corp"
重启 Docker 服务:
sudo systemctl daemon-reload
sudo systemctl restart docker
TensorRT-LLM 配合 Triton Inference Server 为 LLM 推理提供了企业级的高性能解决方案。通过上述步骤,可以实现从模型转换、引擎构建到服务部署的全流程。在实际应用中,需注意模型兼容性、显存管理以及网络配置等问题。对于特定模型(如 Qwen),需针对性调整预处理和后处理逻辑,以确保推理效果稳定。
使用 INT4 或 FP8 量化可以显著减少显存占用并提升推理速度。TensorRT-LLM 支持多种量化方案,如 GPTQ 和 AWQ。选择量化方案时需权衡精度损失与性能收益。
传统的批处理需要等待所有请求完成才能开始下一批,而 Inflight Batching 允许在生成过程中动态插入新请求,提高 GPU 利用率。这要求模型引擎支持 --use_inflight_batching 参数。
Paged KV Cache 解决了传统 KV Cache 内存碎片问题,允许更灵活的序列长度管理。这对于长文本生成场景尤为重要,能有效避免 OOM 错误。
对于超大模型,单卡显存不足时可使用多卡并行。通过设置 --world_size 和 --tp_size 参数,可以将模型切分到多个 GPU 上协同工作。

微信公众号「极客日志」,在微信中扫描左侧二维码关注。展示文案:极客日志 zeeklog
生成新的随机RSA私钥和公钥pem证书。 在线工具,RSA密钥对生成器在线工具,online
基于 Mermaid.js 实时预览流程图、时序图等图表,支持源码编辑与即时渲染。 在线工具,Mermaid 预览与可视化编辑在线工具,online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具,curl 转代码在线工具,online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具,Base64 字符串编码/解码在线工具,online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具,Base64 文件转换器在线工具,online
将 Markdown(GFM)转为 HTML 片段,浏览器内 marked 解析;与 HTML转Markdown 互为补充。 在线工具,Markdown转HTML在线工具,online