TensorRT-LLM加速框架的基本使用
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
