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

news2024/10/7 16:21:33

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

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

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

转为TensorRTLLM格式
转为TensorRTLLM引擎
运行推理

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

输出结果