TensorRT-LLM初探(一)运行llama,以及triton tensorrt llm backend服务化
TensorRT-LLM搭建运行环境以及库、模型运行等流程详解
TensorRT-LLM正式出来有半个月了,一直没有时间玩,周末趁着有时间跑一下。
之前玩内测版的时候就需要cuda-12.x,正式出来仍是需要cuda-12.x,主要是因为tensorr-llm中依赖的CUBIN(二进制代码)是基于cuda12.x编译生成的,想要跑只能更新驱动。
I’ve verified with our CUDA team. A CUBIN built with CUDA 12.x will not load in CUDA 11.x. CUDA 12.x is required to use TensorRT-LLM.
因此,想要快速跑TensorRT-LLM,建议直接将nvidia-driver升级到535.xxx,利用docker跑即可,省去自己折腾环境,至于想要自定义修改源码,也在docker中搞就可以。
理论上替换原始代码中的该部分就可以使用别的cuda版本了(batch manager只是不开源,和cuda版本应该没关系,主要是FMA模块,另外TensorRT-llm依赖的TensorRT有cuda11.x版本,配合inflight_batcher_llm
跑的triton-inference-server也和cuda12.x没有强制依赖关系):
tensorrt-llm中预先编译好的部分
说完环境要求,开始配环境吧!
搭建运行环境以及库
首先拉取镜像,宿主机显卡驱动需要高于等于535:
这个镜像是前几天刚出的,包含了运行TensorRT-LLM的所有环境(TensorRT、mpi、nvcc、nccl库等等),省去自己配环境的烦恼。
拉下来镜像后,启动镜像:
接下来的操作全在这个容器里。
编译tensorrt-llm
首先获取git仓库,因为这个镜像中只有运行需要的lib,模型还是需要自行编译的(因为依赖的TensorRT,用过trt的都知道需要构建engine),所以首先编译tensorrRT-LLM:
然后进入仓库进行编译:
一般不会有环境问题,这个docekr中已经包含了所有需要的包,执行build_wheel
的时候会按照脚本中的步骤pip install一些需要的包,然后运行cmake和make编译文件:
然后pip install tensorrt_llm-0.5.0-py3-none-any.whl
即可。
运行
首先编译模型,因为最近没有下载新模型,还是拿旧的llama做例子。其实吧,其他llm也一样(chatglm、qwen等等),只要trt-llm支持,编译运行方法都一样的,在hugging face下载好要测试的模型即可。
这里我执行:
然后就是TensorRT的编译、构建engine的过程(因为使用了plugin,编译挺快的,这里我只用了一张A4000,所以没有设置world_size,默认为1),这里有很多细节,后续会聊。
编译好engine后,会生成/work/trtModel/llama/1-gpu
,后续会用到。
然后克隆https://github.com/triton-inference-server/tensorrtllm_backend
:
执行以下命令:
然后修改triton_model_repo/中的config:
triton_model_repo/preprocessing/config.pbtxt
Name | Description |
| The path to the tokenizer for the model.这里我改成/work/models/GPT/LLAMA/llama-7b-hf |
| The type of the tokenizer for the model, |
triton_model_repo/tensorrt_llm/config.pbtxt
Name | Description |
| Controls streaming. Decoupled mode must be set to |
| Set to |
| Path to the TensorRT-LLM engines for deployment. In this example, the path should be set to |
triton_model_repo/postprocessing/config.pbtxt
Name | Description |
| The path to the tokenizer for the model. In this example, the path should be set to |
| The type of the tokenizer for the model, |
设置好之后进入tensorrtllm_backend执行:
顺利的话就会输出:
这时也就启动了triton-inference-server,后端就是TensorRT-LLM。
可以看到LLAMA-7B-FP16精度版本,占用显存为:
客户端
然后我们请求一下吧,先走http接口:
triton目前不支持SSE方法,想stream可以使用grpc协议,官方也提供了grpc的方法,首先安装triton客户端:
然后执行:
请求后可以看到是一个token一个token返回的,也就是我们使用chatgpt3.5时,一个字一个字蹦的意思:
因为开了inflight batching,其实可以同时多个请求打过来,修改request_id不要一样就可以:
至此就快速过完整个TensorRT-LLM的运行流程。
使用建议
非常建议使用docker,人生苦短。
在我们实际使用中,vllm在batch较大的场景并不慢,利用率也能打满。TensorRT-LLM和vllm的速度在某些模型上快某些模型上慢,各有优劣。
tensorrt-llm vs vllm
The most fundamental technical difference is that TensorRT-LLM relies on TensorRT ; which is a graph compiler that can produce optimised kernels for your graph. As we continue to improve TensorRT, there will be less and less needs for "manual" interventions to optimise new networks (in terms of kernels as well as taking advantage of numerical optimizations like INT4, INT8 or FP8). I hope it helps a bit.
TensorRT-LLM的特点就是借助TensorRT,TensorRT后续更新越快,支持特性越牛逼,TensorRT-LLM也就越牛逼。灵活性上,我感觉vllm和TensorRT-LLM不分上下,加上大模型的结构其实都差不多,甚至TensorRT-LLM都没有上onnx-parser,在后续更新模型上,python快速搭建模型效率也都差不了多少。
先说这么多,后续会更新些关于TensorRT-LLM和triton相关的文章。
参考
https://github.com/NVIDIA/TensorRT-LLM/issues/45
https://github.com/NVIDIA/TensorRT-LLM/tree/main
https://github.com/NVIDIA/TensorRT-LLM/issues/83