TensorRT-LLM+ChatGLM/Qwen推理部署

一、背景

LLM模型训练完毕、部署在线上进行使用时，遇到的一大难题是并发问题。

以单个模型为例，当模型进行计算时至少要花6、7秒甚至更多。这导致，如果这时有一个新的请求给到模型，模型需要等上一个请求计算完毕才能进行新的计算。

如果一个基于大模型的线上服务，其峰值并发访问量会达到几千、几万，从成本的角度考虑，也不可能同时部署这么多个模型来支持这样的访问量。

另一方面，从理论上来说，LLM模型都是多层transformer block。当前层block计算完毕后，即使最终结果仍为给出，但当前层已处于闲置状态，如果这时把新的请求给到当前层进行计算，也不会对上一个请求的计算结果产生干扰。因此，从理论角度，模型的推理加速是有可行性的。

基于此理论，Nvidia官方则推出了的模型部署方案TensorRT-LLM。

二、部署流程

使用TensorRT-LLM部署大模型需要以下几步：

1. 安装TensorRT-LLM
2. 启动容器
3. 安装python依赖
4. 转化模型格式
5. 运行/部署

下面将依次介绍这些流程具体实施方案。

三、安装TensorRT-LLM

结合自己的显卡驱动版本，从官网拉取合适的docker镜像。
https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tensorrt/tags
我的驱动版本为：Driver Version: 530.30.02 CUDA Version: 12.1
我拉取的版本为：24.03-py3

注意：

1. 早期，官方未提供镜像的时，需要从github拉取源码+cmake本地编译完成安装，这也是其他博客、教程中主要使用的方式。
2. 但这一方式的缺点是需要预装很多环境，且编译过程非常耗时(几个小时)。
   并且，若服务器本身无法连上境外网站(比如github)时，需要先下载到本地再上传至服务器，过程非常繁琐。
3. 因此，强烈推荐大家直接拉取镜像即可。另外，若docker版本低于19.02，其pull命令会略有不同，详情请在官网中查看。

四、启动容器

使用以下指令启动容器

docker run \
  --gpus all \
  --name trt_llm \
  -d \
  --ipc=host \
  --ulimit memlock=-1 \
  --restart=always \
  --ulimit stack=67108864 \
  -v /path/to/your/chatglm_model:/root/chatglm \
  -v /path/to/your/qwen_model:/root/qwen \
  nvcr.io/nvidia/tritonserver:24.03-trtllm-python-py3

这里简单解释一下这里的几个指令：

• –gpus all: 将所有可用的GPU设备暴露给容器，如果只希望容器使用部分gpu，可改为--gpus '"device=1,2"'
• –ulimit memlock=-1: 允许容器锁定任意大小的内存
• –restart=always: 让容器退出时自动重启
• –ulimit stack=67108864: 设置容器内进程的栈大小为64MB
• -v /path/to/your/chatglm_model:/root/chatglm: 将本地模型挂载至容器

容器启动后，使用如下命令进入容器内进行交互操作。

docker exec -it trt_llm /bin/bash 

五、安装依赖

从这一步开始，可直接参考TensorRT-LLM Examples中的示例，根据使用的模型安装相关依赖、运行脚本。

我这里以Chatglm3-6b和Qwen-1.8b-chat为例，在这个examples文件夹下，只需要使用以下几个文件：chatglm/requirements.txt、chatglm/convert_checkpoint.py、qwen/convert_checkpoint.py、run.py和utils.py。

新建相应的文件、把文件内容拷贝进来后，使用pip install -r requirements.txt安装依赖即可。

注意：

1. qwen的requirement.txt比chatglm的多几个依赖项，但在convert_checkpoint.py、run.py、utils.py中并未用上，所以只用chatglm的依赖也可以。
2. chatglm和qwen模型架构不一样，所以相应的convert_checkpoint脚本会稍有不同。
3. qwen和llama架构一致，因此，其他基于llama架构的模型用跟qwen同样的方案即可。

六、转换格式

格式转换分为两部分
(1) 使用convert_checkpoint脚本，将模型从pytorch/tensorflow/paddle等格式统一转换为通用的safetensors格式。
(2) 使用trtllm-build工具，将模型从safetensors格式转换为TensorRT可使用的engine格式。

1. Convert

使用convert_checkpoint脚本，将模型从pytorch格式转换成safetensors格式。

# ChatGLM
python3 convert_checkpoint.py --model_dir /root/chatglm  --output_dir ./tllm_chatglm_fp16_1_gpu --dtype float16
# Qwen
python3 convert_checkpoint.py --model_dir /root/qwen  --output_dir ./tllm_qwen_fp16_1_gpu --dtype float16

2. Build

使用trtllm-build工具，将safetensors格式的模型转换为engine格式

# ChatGLM
trtllm-build --checkpoint_dir ./tllm_chatglm_fp16_1_gpu --output_dir ./trt_engines/chatglm-6b/fp16_1_gpu --gemm_plugin float16
# Qwen
trtllm-build --checkpoint_dir ./tllm_qwen_fp16_1_gpu --output_dir ./trt_engines/qwen-1.8b/fp16_1_gpu --gemm_plugin float16

注意：

1. 此处使用的是单gpu+fp16的转换方式。
2. 要使用多gpu或其他精度，可参考examples/chatglm中的相关部分

七、运行部署

1. 运行测试

以Qwen为例，运行以下脚本即可在终端看到输出

python3 run.py \
  --input_text "你好，你是谁" \
  --max_output_len 50 \
  --tokenizer_dir /root/qwen \
  --engine_dir trt_engines/qwen-1.8b/fp16_1_gpu

结果

2. 部署

部署方式可选FastAPI，也可以选择使用官方的Nvidia-Triton推理框架。

八、总结

使用TensorRT-LLM部署大模型包括以下几个步骤：安装TensorRT-LLM、启动容器、安装python依赖、转化模型格式、部署服务。在官方对这一框架提供了较为完善的支持后，以上这些步骤已经可以很方便的完成了。

九、参考文档

除了官方的文档、代码外，还有一些博客、视频对我帮助很大，在此也特别感谢这些作者的分享！

TensorRT源码: https://github.com/NVIDIA/TensorRT-LLM
TensorRT镜像: https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tensorrt/tags
Triton框架: https://www.nvidia.cn/ai-data-science/products/triton-inference-server

博客: docker编译TensorRT-LLM镜像 这位博主写了一系列关于TensorRT的文章，非常推荐大家都看看
视频: TensorRT-LLM实战指南 作者有一整套TensorRT实战视频，涉及方面也很全面，推荐大家观看

后记：
当我写这篇博客的时候，我参考的这篇博客作者换了他知乎的头像，然后我发现他跟那个视频的作者居然就是同一个人！(巧了嘛这不是)