基于TensorRT-LLM的Yuan 2.0推理服务部署

自Yuan2.0大模型发布以来，团队一直致力于提高大模型的性能和效率。最近我们完成了基于TensorRT-LLM推理框架的Yuan2.0大模型部署实践，一起来尝试一下吧！

TensorRT-LLM是2023年10月新发布的大模型推理部署框架，用其部署Yuan2.0大模型一方面可以显著提高模型推理速度和减少延时，让用户可以更快的获取推理结果；另一方面，利用TensorRT-LLM灵活的插件功能，使得用户在部署Yuan2.0时更有效的利用硬件资源从而节省硬件成本。此外， TensorRT-LLM构建的Yuan2.0 engine可以进一步使用Nvidia提供的Tritonserver（tensorrtllm_backend），生成易于集成的接口和文档，使得用户能够轻松地将其集成到现有的深度学习推理系统中并快速实现业务需求。

 研发环境要求：

快速导航

 一、配置TensorRT-LLM和tensorrtllm_backend环境

二、下载Yuan2.0-2B权重文件，克隆Github Yuan2.0仓库

三、构建Yuan2.0 trt_llm_engine

四、启动Tritonserver服务 

 五、性能测试结果和局限

---

 一、配置TensorRT-LLM和tensorrtllm_backend环境

为了简化TensorRT-LLM部署Yuan2.0模型过程，我们提供了一个镜像。其中包含了创建、运行TensorRT-LLM模型和启动 Tritonserver所需的一切环境和依赖。环境配置按照以下步骤操作即可：

•  拉取我们提供的镜像，参考地址：

https://hub.docker.com/repository/docker/zhaoxudong01/trt_llm_yuan/tags?page=1&ordering=last_updated

docker pull zhaoxudong01/trt_llm_yuan:v1.0

•  根据上述镜像创建容器

docker run --gpus all -itd --network=host --cap-add=IPC_LOCK --device=/devinfiniband --privileged --name TensorRT-LLM-Yuan --ulimit core=0 --ulimit memlock=1 --ulimit stack=68719476736 --shm-size=1000G zhaoxudong01/trt_llm_yuan:v1.0

• 进入容器

docker exec -it TensorRT-LLM-Yuan bash

关于镜像我们还提供了百度网盘的下载地址：

TensorRT-LLM yuan2.0部署镜像 I 百度网盘

下载好镜像文件后，通过下面命令加载镜像

docker load -i zhaoxudong01_trt_llm_yuan_v1.0.tar

通过下面命令可以查看是否加载镜像成功

docker ps -a

二、下载Yuan2.0-2B权重文件，克隆Github Yuan2.0仓库

 Github上Yuan2.0的仓库除了包含Yuan2.0的开源代码外，也包含了我们对Yuan2.0-2B适配TensorRT-LLM和tensorrtllm_backend的代码，二者都在Yuan-2.0/3rdparty/目录下。Yuan2.0的权重文件和github仓库下载步骤如下：

• 下载Yuan2.0-2B的权重文件

 Yuan2.0-2B的hugging face模型下载地址：https://huggingface.co/IEITYuan/Yuan2-2B-hf

如果外网访问困难，也可以通过魔搭社区、wisemodel、百度网盘、OpenXlab等平台进行下载。方便起见，我们将不同规模的模型下载链接附上，后续的部署过程同Yuan2-2B-Janus-hf模型完全相同。

模型下载完成后解压，文件夹中内容如下图所示，即表示模型下载完整了 

 

•  将下载好的文件上传至刚刚创建的docker的/mnt/Yuan2.0-2B目录下

cd /mnt/  #进入/mnt目录下
mkdir Yuan2.0-2B  #在mnt目录下创建Yuan2.0-2B文件夹
cd Yuan2.0-2B  #进入/mnt/Yuan2.0-2B目录，将刚刚下载的权重文件拷贝至此文件夹中

•  从github上克隆Yuan2.0仓库

cd /mnt  #回到/mnt目录下
git clone https://github.com/IEIT-Yuan/Yuan-2.0.git  #克隆Yuan2.0项目仓库

三、构建Yuan2.0 trt_llm_engine

TensorRT-LLM构建Yuan2.0 trt_llm_engine的脚本和测试engine的脚本分别在Yuan-2.0/3rdparty/TensorRT-LLM/examples/yuan下的build.py和run.py，构建和测试engine的步骤如下：

• 进入克隆到本的Yuan2.0代码仓库下的TensorRT-LLM目录

cd /mnt/Yuan-2.0/3rdparty/TensorRT-LLM

• 构建Yuan2.0-2B的TensorRT-LLM引擎(称为Yuan2B_trt_llm_engine) 

cd examples/yuan  #进入TensorRT-LLM下yuan的目录
python3 build.py --max_batch_size=4 --model_dir=/mnt/Yuan2.0-2B --dtype=float16  
--use_gemm_plugin=float16 --output_dir=/mnt/Yuan2.0-2B/trt_engines/fp16/1-gpu  
#其中‘--model_dir’是模型权重文件的路径，‘--output_dir’是生成的trt_llm_engine的路径

这里需要记录--output_dir指定的路径，它是生成的trt_llm_engine的路径，在后续启动tritionserver时，配置config.pbtxt需要用到。

•  测试Yuan2B_trt_llm_engine

python3 run.py --max_output_len=100 --tokenizer_dir=/mnt/Yuan2.0-2B --output_dir=/mnt/Yuan2.0-2B/trt_engines/fp16/1-gpu --input_text="写一篇春游作文<sep>"  
#用run.py测试生成的engine，其中‘--max_output_len’ 是最大输出长度、
#‘--tokenizer_dir’是tokenizer的路径、
#‘--output_dir’是engine的路径、
#‘—input_text’是提示词(格式是：“提示内容+<sep>”)

•  如果Yuan2B_trt_llm_engine构建成功，run.py的预期输出结果如下:

Input [Text 0]: "写一篇春游作文<sep>"
Output [Text 0 Beam 0]: "亲爱的同学们：
春天到了，大地复苏，草木变绿，花儿绽放。这是一个美好的季节，也是我们展开青春梦想的舞台。今天，我们前往公园，感受春天的温暖。
早晨，阳光洒在大地上，给空气涂抹了一层斑斓的彩色，缤纷的色彩像云霞一样美丽。鸟儿开始欢快地舞蹈，漫过天空，唤醒大地。我们沐浴在温暖的阳光下，感受到来自大自然的恩赐。
漫步在小径上，我们"

由于模型生成内容有随机性，output的内容可能不一样。受限于输出序列长度，输出的内容最终也不一定完整，可通过循环调用和限制截断符等后处理操作实现输出标准化。这里仅将模型转成tensorRT实现推理加速，不涉及后处理步骤。

四、启动Tritonserver服务 

 当成功构建Yuan2B_trt_llm_engine后，可以利用Tritonserver去启动Yuan2.0的推理服务。推理服务的代码在Yuan-2.0/3rdparty/tensorrtllm_backend下，启动推理服务的步骤如下：

• 进入tensorrtllm_backend目录

cd /mnt/Yuan-2.0/3rdparty/tensorrtllm_backend

• 拷贝模型文件

cp -r all_models/gpt triton_model_repo_yuan

• 修改tensorrtllm_backend/triton_model_repo_yuan下模型参数

'triton_mode_repo_yuan' 目录下有四个模型将被使用：

1. "preprocessing": 此模型用于进行编码，即从提示（字符串）转换为 input_ids（整数列表）
2. "tensorrt_llm": 此模型是您的 TensorRT-LLM 模型的封装，用于推理。
3. "postprocessing": 此模型用于解码，即将输出 ID（整数列表）转换为输出字符串。
4. "ensemble": 此模型用于将上述三个模型连接在一起 preprocessing -> tensorrt_llm -> postprocessing

其中，部署模型需要修改的文件主要在preprocessing、tensorrt_llm和postprocessing目录下。

1.修改preprocessing下的config.pbtxt

#1.修改preprocessing下的config.pbtxt
vim /mnt/Yuan-2.0/3rdparty/tensorrtllm_backend/triton_model_repo_yuan/preprocessing/config.pbtxt

按照下图修改line62和line69

2.修改tensorrtllm下的config.pbtxt 

#2.修改tensorrtllm下的config.pbtxt
vim /mnt/Yuan-2.0/3rdparty/tensorrtllm_backend/triton_model_repo_yuan/tensorrtllm/config.pbtxt

按照下图修改line131

 3.修改postprocessing下的config.pbtxt

#3.修改postprocessing下的config.pbtxt
vim /mnt/Yuan-2.0/3rdparty/tensorrtllm_backend/triton_model_repo_yuan/postprocessing/config.pbtxt

按照下图修改line22

•  启动Tritonserver服务

回到tensorrtllm_backend目录下，启动Tritonserver服务：

cd /mnt/Yuan-2.0/3rdparty/tensorrtllm_backend
#发起Tritonserver服务
python3 scripts/launch_triton_server.py --world_size=1 --model_repo=triton_model_repo_yuan

成功发起服务后，服务器会产生类似以下的日志：

I0919 14:52:10.475738 293 grpc_server.cc:2451] Started GRPCInferenceService at 0.0.0.0:8001
I0919 14:52:10.475968 293 http_server.cc:3558] Started HTTPService at 0.0.0.0:8000
I0919 14:52:10.517138 293 http_server.cc:187] Started Metrics Service at 0.0.0.0:8002

• 使用 Triton 的 generate 端点查询服务器

从 Triton 23.10 版本开始，您可以使用 Triton 的generate endpoint，通过以下常规格式的 curl 命令在客户端环境/容器中查询服务器：

curl -X POST localhost:8000/v2/models/ensemble/generate -d '{"text_input": "写一篇春游作文<sep>", "max_tokens": 100, "bad_words": "", "stop_words": "", "pad_id": 77185, "end_id": 77185}'

输出结果如下：

ok，到这里我们就完成了基于TensorRT-LLM的yuan2.0大模型推理服务部署。

•  任务结束后停止Tritonserver服务

pkill tritonserver

 五、性能测试结果和局限

• 性能测试结果

 我们比较了Yuan2.0-2B的trt_llm模型和原始的megatron模型进行的推理速度

max_output_len=300, prompt="写一篇春游作文"

•  局限

到目前为止，使用TensorRT-LLM构建的trt-llm-Yuan2.0引擎时：

• 暂不支持'gpt_attention'，'remove_input_padding'和'inflight_batching'
• 只支持Yuan2.0-2B单GPU部署

如需获取关于Yuan2.0和TensorRT-LLM的信息，请参考：

Yuan2.0: GitHub - IEIT-Yuan/Yuan-2.0: Yuan 2.0 Large Language Model

TensorRT-LLM: GitHub - NVIDIA/TensorRT-LLM: TensorRT-LLM provides users with an easy-to-use Python API to define Large Language Models (LLMs) and build TensorRT engines that contain state-of-the-art optimizations to perform inference efficiently on NVIDIA GPUs. TensorRT-LLM also contains components to create Python and C++ runtimes that execute those TensorRT engines.