1. TensorRT-LLM 编译模型
1.1 TensorRT-LLM 简介
使用 TensorRT 时,通常需要将模型转换为 ONNX 格式,再将 ONNX 转换为 TensorRT 格式,然后在 TensorRT、Triton Server 中进行推理。
但这个转换过程并不简单,经常会遇到各种报错,需要对模型结构、平台算子有一定的掌握,具备转换和调试能力。而 TensorRT-LLM 的目标就是降低这一过程的复杂度,让大模型更容易跑在 TensorRT 引擎上。
需要注意的是,TensorRT 针对的是具体硬件,不同的 GPU 型号需要编译不同的 TensorRT 格式模型。这与 ONNX 模型格式的通用性定位显著不同。
同时,TensortRT-LLM 并不支持全部 GPU 型号,仅支持 H100、L40S、A100、A30、V100 等显卡。
1.2 配置编译环境
--gpus device=0 表示使用编号为 0 的 GPU 卡,这里的 hubimage/nvidia-tensorrt-llm:v0.7.1 对应的就是 TensorRT-LLM v0.7.1 的 Release 版本。
由于自行打镜像非常麻烦,这里提供几个可选版本的镜像:
- hubimage/nvidia-tensorrt-llm:v0.7.1
- hubimage/nvidia-tensorrt-llm:v0.7.0
- hubimage/nvidia-tensorrt-llm:v0.6.1
1.3 编译生成 TensorRT 格式模型
在上述容器环境下,执行命令:
生成的文件主要有三个:
- baichuan_float16_tp1_rank0.engine,嵌入权重的模型计算图文件
- config.json,模型结构、精度、插件等详细配置信息文件
- model.cache,编译缓存文件,可以加速后续编译速度
1.4 推理测试
1.5 验证是否严重退化
模型推理优化,可以替换算子、量化、裁剪反向传播等手段,但有一个基本线一定要达到,那就是模型不能退化很多。
在精度损失可接受的范围内,模型的推理优化才有意义。TensorRT-LLM 项目提供的 summarize.py 可以跑一些测试,给模型打分,rouge1、rouge2 和 rougeLsum 是用于评价文本生成质量的指标,可以用于评估模型推理质量。
- 获取原格式模型的 Rouge 指标
由于目前 optimum 不支持 Baichuan 模型,因此,需要编辑 examples/summarize.py 注释掉 model.to_bettertransformer(),这个问题在最新的 TensorRT-LLM 代码中已经解决,我使用的是当前最新的 Release 版本 (v0.7.1)。
输出结果:
- 获取 TensorRT 格式模型的 Rouge 指标
输出结果:
TensorRT-LLM 编译之后的模型,rougeLsum 从 24 降到了 22,说明能力会有退化,但只要在可接受的范围之内,还是可以使用的,因为推理速度会有较大的提升。
完成这步之后,就可以退出容器了,推理是在另外一个容器中进行。
2. Triton Server 配置说明
2.1 Triton Server 简介
Triton Server 是一个推理框架,提供用户规模化进行推理的能力。具体包括:
- 支持多种后端,tensorrt、onnxruntime、pytorch、python、vllm、tensorrtllm 等,还可以自定义后端,只需要相应的 shared library 即可。
- 对外提供 HTTP、GRPC 接口
- batch 能力,支持批量进行推理,而开启 Dynamic batching 之后,多个 batch 可以合并之后同时进行推理,实现更高吞吐量
- pipeline 能力,一个 Triton Server 可以同时推理多个模型,并且模型之间可以进行编排,支持 Concurrent Model Execution 流水线并行推理
- 观测能力,提供有 Metrics 可以实时监控推理的各种指标
图片
上面是 Triton Server 的架构图,简单点说 Triton Server 是一个端(模型)到端(应用)的推理框架,提供了围绕推理的生命周期过程管理,配置好模型之后,就能直接对应用层提供服务。
2.2 Triton Server 使用配置
在 Triton 社区的示例中,通常会有这样四个目录:
对于 Triton Server 来说,上面的目录格式实际上是定义了四个模型,分别是 preprocessing、tensorrt_llm、postprocessing、ensemble,只不过 ensemble 是一个组合模型,定义多个模型来融合。
ensemble 存在的原因在于 tensorrt_llm 的推理并不是 text2text ,借助 Triton Server 的 Pipeline 能力,通过 preprocessing 对输入进行 Tokenizing,postprocessing 对输出进行 Detokenizing,就能够实现端到端的推理能力。否则,在客户端直接使用 TensorRT-LLM 时,还需要自行处理词与索引的双向映射。
这四个模型具体作用如下:
- preprocessing, 用于输入文本的预处理,包括分词、词向量化等,实现类似 text2vec 的预处理。
- tensorrt_llm, 用于 TensorRT 格式模型的 vec2vec 的推理
- postprocessing,用于输出文本的后处理,包括生成文本的后处理,如对齐、截断等,实现类似 vec2text 的后处理。
- ensemble,将上面的是三个模型进行融合,提供 text2text 的推理
上面定义的模型都有一个 1 目录表示版本 1 ,在版本目录中放置模型文件,在模型目录下放置 config.pbtxt 描述推理的参数 input、output、version 等。
2.3 模型加载的控制管理
Triton Server 通过参数 --model-control-mode 来控制模型加载的方式,目前有三种加载模式:
- none,加载目录下的全部模型
- explicit,加载目录下的指定模型,通过参数 --load-model 加载指定的模型
- poll,定时轮询加载目录下的全部模型,通过参数 --repository-poll-secs 配置轮询周期
2.4 模型版本的控制管理
Triton Server 在模型的配置文件 config.pbtxt 中提供有 Version Policy,每个模型可以有多个版本共存。默认使用版本号为 1 的模型,目前有三种版本策略:
- 所有版本同时使用
- 只使用最近 n 个版本
- 只使用指定的版本
3. Triton Server 中使用 TensorRT-LLM
3.1 克隆配置文件
本文示例相关的配置已经整理了一份到 GitHub 上,拷贝模型到指定的目之后,就可以直接进行推理了。
3.2 组织推理目录
- 拷贝 TensorRT 格式模型
- 拷贝源模型
此时文件的目录结构是:
3.3 启动推理服务
如果一台机器上运行了多个 triton server,那么需要用 shm-region-prefix-name=prefix0_ 区分一下共享内存的前缀,详情可以参考 https://github.com/triton-inference-server/server/issues/4145 。
启动日志:
四个模型都处于 READY 状态,就可以正常推理了。
- 查看模型配置参数
可以查看模型的推理参数。如果使用的是 auto-complete-config,那么这个接口可以用于导出 Triton Server 自动生成的模型推理参数,用于修改和调试。
- 查看 Triton 是否正常运行
3.4 客户端调用
- 安装依赖
Triton GRPC 接口的性能显著高于 HTTP 接口,同时在容器中,我也没有找到 HTTP 接口的示例,这里就直接用 GRPC 了。
- 推理测试
3.5 查看指标
Triton Server 已经提供了推理指标,监听在 8002 端口。在本文的示例中,就是 38002 端口。
在 Grafana 中可以导入面板 https://grafana.com/grafana/dashboards/18737-triton-inference-server/ 查看指标,如下图:
图片
4. 总结
本文主要是在学习使用 TensorRT 和 Triton Server 进行推理过程的记录,主要内容如下:
- TensorRT 是一种针对 Nvidia GPU 硬件更高效的模型推理引擎
- TensorRT-LLM 能让大模型更快使用上 TensorRT 引擎
- Triton Server 是一个端到端的推理框架,支持大部分的模型框架,能帮助用户快速实现规模化的推理服务
- Triton Server 下使用 TensorRT-LLM 进行推理的示例
5. 参考
- https://mmdeploy.readthedocs.io/zh-cn/latest/tutorial/03_pytorch2onnx.html
- https://docs.nvidia.com/deeplearning/tensorrt/container-release-notes/running.html#running
- https://github.com/NVIDIA/TensorRT-LLM
- https://github.com/triton-inference-server/triton-tensorrtllm
- https://zhuanlan.zhihu.com/p/663748373