TensorRT-LLM是英伟达发布的针对大模型的加速框架，TensorRT-LLM是TensorRT的延申。TensorRT-LLM的GitHub地址是 https://github.com/NVIDIA/TensorRT-LLM

这个框架在0.8版本有一个比较大的更新，原先的逻辑被统一了，所以早期的版本就不介绍了。这个教程介绍0.11.0的main分支版本。

TensorRT-LLM的加速推理分为3个阶段。：

1. 转为TensorRTLLM格式
2. 转为TensorRTLLM引擎
3. 运行推理

TensorRT-LLM只支持部分大模型，具体可以看它GitHub中的example文件夹。下面以InternLM2大模型为例子，展示TensorRT-LLM的使用。

环境

• 8张3090，在本例中使用了两张24G的3090显卡。
• 现在基本都使用Docker了，因为使用TensorRT-LLM主要还是想要最终的部署，建议参考tensorrtllm_backend服务端
  来编译。这个仓库提供了tritonserver+TensorRTLLM（服务端+加速框架），也有编译好的docker镜像
  NGC地址

转为TensorRT-LLM格式

在官方给出的Internlm2实例文件夹中执行命令

python convert_checkpoint.py --model_dir /tensorrtllm_backend/triton_model_repo/tensorrt_llm/internlm2-chat-20b \--dtype bfloat16 \--output_dir ./internlm2-chat-20b/trt_engines/bf16/2-gpu/ \--tp_size 2

因为一张24G的显卡不足以加载20B的模型 ，所以设置tp_size为2。

生成引擎文件

使用trtllm-build命令将TensorRT-LLM格式的文件生成引擎。命令如下

CUDA_VISIBLE_DEVICES=1,2 trtllm-build --checkpoint_dir  /model2engine/internlm2/internlm2-chat-20b/trt_engines/bf16/2-gpu \--max_batch_size 8  --max_input_len 2048 \--max_output_len 2048  \--gemm_plugin bfloat16\--gpt_attention_plugin bfloat16 \--remove_input_padding enable  \--paged_kv_cache enable \--tp_size 2  \--use_custom_all_reduce disable\--output_dir  /model2engine/internlm2/engine_outputs

运行推理

CUDA_VISIBLE_DEVICES=1,2  mpirun -n 2 --allow-run-as-root \python ../run.py --max_output_len=120 \--input_text '你是谁啊？' \--tokenizer_dir /tensorrtllm_backend/triton_model_repo/tensorrt_llm/internlm2-chat-20b \--engine_dir=/model2engine/internlm2/engine_outputs