TensorRT-LLM加速框架的基本使用

TensorRT-LLM是英伟达发布的针对大模型的加速框架，TensorRT-LLM是TensorRT在大模型时代的延伸。TensorRT-LLM的GitHub地址是 https://github.com/NVIDIA/TensorRT-LLM。相比于其他的加速框架，TensorRT-LLM是我们用过的最快的。

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

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

将大模型转为TensorRT-LLM格式
将TensorRT-LLM格式转为TensorRT-LLM引擎
运行推理

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

环境

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

转为TensorRT-LLM格式

首先进入docker镜像中，在官方给出的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

其中model_dir是从Huggingface下载的大模型权重地址，output_dir 为输出的TensorRT-LLM格式的文件地址。因为一张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

转化完成

运行推理

在example目录中有run.py文件，该文件可以快速测试引擎文件是否有问题。

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

输出结果

对于TensorRT-LLM的建议

目前TensorRTLLM对新模型的支持有点慢，有的模型需要在发布几个月后才会在正式版中受到支持。
NGC镜像更新慢，得等上十天半个月才会有更新，然而手动编译的镜像体积又太大。希望官方可以出个说明解释下NGC镜像是怎么编译的。