【大模型实战篇】利用TensorRT LLM部署及推理大模型实战（bloom560m、chatglm-6b等大模型）

1. TensorRT-LLM基础介绍

        TensorRT-LLM 是 NVIDIA 用于做 LLM（Large Language Model）的可扩展推理方案。基于 TensorRT 深度学习编译框架来构建、编译并执行计算图，并借鉴了许多 FastTransformer 中高效的 Kernels 实现，然后利用 NCCL 完成设备之间的通讯。TensorRT-LLM 提供了一个易于使用的 Python API，让用户能够定义大型语言模型（LLMs）并构建包含 TensorRT 引擎，以便在 NVIDIA GPU 上执行推理【1，2，3】。

     【3，4】对TensorRT-LLM的架构进行了介绍。TensorRT + FastTransFormer + NCCL + Python api => TensorRT-LLM。

        传统TensorRT需要将Pytorch模型转成Onnx，然后再将Onnx转成TensorRT。一般还需要做数据对齐，plugin编写，Onnx修改来适配TensorRT plugin。深度学习模型通常使用各种框架（如PyTorch、TensorFlow、Keras等）进行训练和部署，而每个框架都有自己的模型表示和存储格式。ONNX（Open Neural Network Exchange） 是一个开放的、跨平台的神经网络模型交换格式。主要目标是实现机器学习模型的互操作性，使得模型可以在不同的框架和运行时之间轻松迁移。ONNX 模型由节点和边组成的计算图表示。ONNX使用Protobuf作为其模型文件的序列化格式。Protobuf是一种轻量级的、高效的数据交换格式，但它在序列化和反序列化大型数据时有一个默认的大小限制。在Protobuf中，默认的大小限制是2GB。单个序列化的消息不能超过2GB的大小。

        TensorRT-LLM 作为优化大语言模型（LLM）推理的库，提供了包括自定义注意力核、批处理、分页键值缓存、量化（如 FP8、INT4 AWQ、INT8 SmoothQuant 等）以及更多优化手段，在 NVIDIA GPU 上高效地进行推理。提供了一个 Python API，用于将 LLM 构建为优化过的 TensorRT 引擎。包含 Python 绑定和 C++ 运行时，用于执行这些 TensorRT 引擎。此外，它还包括与 NVIDIA Triton 推理服务器集成的后端。使用 TensorRT-LLM 构建的模型可以在各种配置下运行，从单个 GPU 到多个节点、多个 GPU（使用张量并行和/或流水线并行）。此外，也提供了几种预定义的流行模型。通过类似 PyTorch 的 Python API修改和扩展以满足自定义需求。

        TensorRT-LLM 构建于 TensorRT 深度学习推理库之上。它利用了 TensorRT 大部分的深度学习优化，并在此基础上添加了 LLM 特定的优化，TensorRT 是一个预编译器，它构建“引擎”，这些引擎是编译后的模型的优化表示，包含整个执行图。这些引擎是针对特定的 GPU 架构进行优化的，可以在生产环境中进行验证、基准测试和序列化以便后续部署。

2. TensorRT-LLM部署大模型推理实战

        实战一般绕不开环境配置，tensorrt-llm也是如此，一开始我是按照官网提供的案例【5】来操纵，不出意外得出了各种报错，要么是系统依赖的版本问题，要么是编译mpi出错，搞了半天。后来实在觉得这条路子不好走，就换了一种思路。

        （1）先去下载nvidia的NVIDIA Container Toolkit，绕过手动配置环境的困扰。这里docker pull镜像的时候，最好加一下国内的加速代理，不然又要很崩溃。当然也可以在本地用魔法先下载镜像，然后采用docker save -o  xxx.tar  #对应container保存镜像，然后上传到服务器，再执行docker load -i xxx.tar，然后再启动容器。

        docker镜像地址：https://hub.docker.com/r/nvidia/cuda/tags?page=2&name=12.4

docker pull nvidia/cuda:12.4.0-devel-ubuntu22.04 

docker run --rm --ipc=host --runtime=nvidia --gpus all --entrypoint /bin/bash -it nvidia/cuda:12.4.0-devel-ubuntu22.04

        成功启动docker镜像后，输入docker ps找到对应容器，进入容器环境：

 docker exec -it xxx bash

        这里可能需要注意下，如果本地显卡的cuda版本最好和docker镜像的配置一致，不然后面又会各种报错。还有建议下载ubuntu22.04系统的版本，带python3.10版本。

        （2）接下来在容器中安装TensorRT-LLM，这一步需要花些时间。

# Install dependencies, TensorRT-LLM requires Python 3.10
apt-get update && apt-get -y install python3.10 python3-pip openmpi-bin libopenmpi-dev git git-lfs

pip3 install tensorrt_llm -U --extra-index-url https://pypi.nvidia.com

# Check installation
python3 -c "import tensorrt_llm"        

2.1 部署运行bloom560m

        需要git clone TensorRT-LLM的项目，并且安装对应大模型的依赖。

git clone https://github.com/NVIDIA/TensorRT-LLM.git

cd TensorRT-LLM

pip install -r examples/bloom/requirements.txt

git lfs install

        按照惯例，我们从modelscope上拉取对应大模型的参数文件。

git clone https://modelscope.cn/models/AI-ModelScope/bloom-560m.git

        然后执行三部曲：（1）执行模型参数convert；（2）build tensorrt engine；（3）运行engine

        从推理结果看，bloom-560m的效果貌似很差。

2.2 部署运行chatglm-6b

        从modelscope上拉取对应大模型的参数文件。

（1）执行convert_checkpoint.py  

cd ~/TensorRT-LLM/examples/chatglm   

python3 convert_checkpoint.py --model_dir /root/chatglm3-6b/ \
                --dtype float16 \
                --output_dir /root/TensorRT-LLM/examples/chatglm/6b/trt_ckpt/fp16/1-gpu/

（2）build tensorrt engine.

trtllm-build --checkpoint_dir /root/TensorRT-LLM/examples/chatglm/6b/trt_ckpt/fp16/1-gpu/ \
                --gemm_plugin float16 \
                --output_dir /root/TensorRT-LLM/examples/chatglm/6b/trt_engines/fp16/1-gpu/

（3）运行engine

python3 /root/TensorRT-LLM/examples/run.py --input_text "What's result of 1 plus 1?" \
                  --max_output_len 50 \
                  --tokenizer_dir /root/chatglm3-6b/  \
                  --engine_dir /root/TensorRT-LLM/examples/chatglm/6b/trt_engines/fp16/1-gpu/

        从结果看，chatglm-6b的推理结果还可以。

3. 参考材料

【1】TensorRT-LLM部署Qwen

【2】TensorRT-LLM实战

【3】​NVIDIA大模型推理框架TensorRT-LLM

【4】TensorRT-LLM实战指南

【5】TensorRT-LLM/examples/qwen/README.md at main · NVIDIA/TensorRT-LLM · GitHub