TensorRT-LLM初探（三）最佳部署实践

本篇文章根据NVIDIA AI技术开放日 2024 夏[1]中TRT-LLM 最佳部署实践[2]的演讲，结合自己的一些经验整理成本篇文章，算是TensorRT-LLM初探第三篇——最佳部署实践。

下文图片PPT部分皆来源于TRT-LLM 最佳部署实践[3]。

之前两篇的传送门：

TensorRT-LLM初探（一）基于最新commit运行llama，以及triton-tensorrt-llm-backend[4]
TensorRT-LLM初探（二）简析了结构，用的更明白[5]

本篇根据讲座的内容也大概分为以下几点：

TensorRT-LLM介绍
端到端workflow
如何debug
如何添加新的模型

TRT-LLM简单再介绍

TensorRT-LLM的介绍前几篇中已提到，就不过多赘述了。
这里列一个TensorRT-LLM的功能和定位：

TRT-LLM和vllm、lmdeploy、sglang[6]一样，提供大模型的推理支持，包含了大模型推理的：

模型结构，提前定义好的模型结构
runtime调度（inflight batching、kv cache reuse）
kernels（MMHA、FMHA）
量化技术（FP8、INT8、INT4、kv cache）

这里挨个过下：

模型结构

模型结构就是提前定义好的llama或者其他大模型的网络结构，直接复用就行。

搭建好的模型可以使用TensorRT帮你生成kernel，和小模型走onnx的路子不一样，trt-llm完善了TensorRT-python-api，使其更好用和易于搭建，更灵活一点，不过说实话，相比使用vllm搭建还是稍微难一点。

kernel优化

对于大模型来说，简单对于kernel的优化是不够的。之前小模型的经验，优化模型第一直觉就是优化kernel，但是对于大模型来说runtime、调度也很重要。

优化kernel直接可以优化模型性能，降低latency；而runtime或者说调度可以提升整体的吞吐。

目前trt-llm中比较常用的就是MMHA（MaskedMultiheadAttention）和FMHA（FusedMultiheadAttention），这俩都是fused Multi Head Attention，MMHA是context那边的MMHA，也是fused的，这俩都是从faster transformer借鉴来，目前也更新了很多版本。

runtime调度

Runtime调度也很重要，除了最基本的inflight batching之外，kv cache优化目前更重要一些。

因为目前长context的需求现在比较多，缓存kv cache的部分需要的memory越来越重。所以有了kv cache压缩，也就是量化以及low rank的需求。当context变超长的时候，kv cache需要的显存大小甚至会比模型本身还大，所以kv cache压缩也比较重要，比如INT4。

Runtime方式变化也会影响kernel实现的方式，需要修改kernel的实现方式去配合runtime。另外，TP和PP也是标配，算是runtime的一部分，相比于TensorRT只支持单卡，trt-llm增加了多卡的支持，是通过trt-plugin支持的。

量化

最后是量化，量化的支持trt-llm支持也是不少，这里暂时略过。两量化需要单独的篇幅去说，开放日也有单独的讲座去讲量化相关。

关于开源

TRT比较被诟病的就是开源开的不彻底，能改动的地方不多。

究其原因，是有很多针对不同硬件做了定制化优化的kernel，如果放出来，我们就可以通过代码反推到硬件的底层逻辑的设计。

比如FMHA的代码，某一个配置中，可以看到有很多针对不同sm架构的实现代码：

针对不同size不同显卡架构都有不同的实现，相对比较细致、比较极致，总之就是和硬件比较相关的代码没有开源，其他的代码开源。

值得一提是还有一个runtime代码，之前GptSession好歹是开源的，后来切换成了Executor，直接给你闭源了，想要添加功能只能去官方提需求，没办法自己修改：

端到端 workflow

这里我们讲下端到端使用trt-llm开发的流程。

首先提一下Python API，相比于使用原始TensorRT的python-api，TRT-LLM在上面又封装了一层，尽可能和pytorch的风格一致，易于搭建新的网络。

不过要注意，python-api只是搭建，搭建好的网络只是TensorRT的网络格式（类似于onnx2trt的中间IR形式），不能直接运行（这点要区分于Pytorch），需要build engine后才可以。实际运行还是使用C++，和TensorRT一样。

TRT-LLM也提供了High Level API，类似于torch和huggingface的关系：

当然流程还是先转换权重格式、搭建trt-llm网络结构、build engine，然后就可以运行了。

回到端到端的workflow，就是上述聊到的那几个流程：

我们首先需要convert checkpoint，将原始的weight转换为trt-llm的格式
然后将转换好的weight全部填入提前定义好的网络结构中
最后build已经读取了权重而且定义好网络结构的network，这里区别于pytorch，因为pytorch是动态图，而trt-llm是静态图，所以相对来说没有那么方便，需要先定义好网络结构，再build才能得到最终的engine结构，这也是优化后的计算图

最后编译出来的engine可以在python中先进行测试，测试没问题后，就可以部署到C++中，最终通过triton上线。

安装 && install

开放日中也简单提了下安装过程，可能是TRT-LLM安装坑确实比较多吧...

第一个是在利用docker自行编译trt-llm源码，也是我比较常用的方式，优点是可以修改源码以及不需要考虑环境，不好的就是对网络要求稍高点（懂得都懂）。

第二种方式是直接通过pip安装，这个建议在之前已经有环境可以跑起来trt-llm的基础上，你想要更新版本，可以这么搞。或者说你有纯净的trt-llm依赖的环境（比如从ngc拉下来的镜像，或者第一个build出来的镜像），直接在这个环境中pip install即可。

如果你想直接在其他开发环境中pip install，可能会和你本地的一些库有一些不兼容的地方（比如你的torch是自己编译的，gcc版本不不一样），可能有些symbolic找不到，所以最好是纯净的环境。

第三种是借用NGC中提供的镜像。

NVIDIA NGC（NVIDIA GPU Cloud）是一个为深度学习、机器学习和高性能计算（HPC）提供优化的GPU软件的中心。这个平台提供了容器、预训练模型、模型脚本和行业解决方案，帮助数据科学家、开发者和研究人员更快地构建解决方案，我们快速开发使用的镜像一般来源于这里

NGC中的镜像已经提前预装了trtllm-triton-backend和trt-llm这俩库，所以trt-llm需要的系统环境也有了。虽然说有预先编译好的trt-llm，其实后续我们也可以自行编译其他版本的，都比较灵活。

最后总结了下各种安装方式的优点和缺点：

转换权重（checkpoint）

之前TRT-LLM每个模型都有一个convert脚本，会比较乱而且不好维护，所以现在TRT-LLM统一了convert接口：

在convert checkpoint的地方统一了之后会有很多好处：

权重转换后需要把权重塞到模型中，需要定义模型结构，trt-llm预先提供了一些比较火的模型结构，对这些个模型提供支持：

第四步就是在权重转换为trt-llm之后，开始进行build构建。有个细节是，在build的时候有很多参数会影响性能，官方预设的参数默认是效果比较好的，但是我们肯定要根据自己实际的需求去调节参数，不论是速度还是精度问题：

在构建好engine之后，就可以开始运行了，建议首先使用run.py在python端进行测试。然后也可以使用其他的.py文件或者gptManagerBenchmark去评测模型精度或者性能：

MMLU、公开的LLM测试集，来测试trt-llm模型build之后的精度，一般就是测试一个pytorch的再测试一个trt-llm的，简单对比即可。

TRT-LLM也提供了benchmark工具，gptManagerBenchmark是提前编译好的可执行文件，专门用来测试性能，也可以测试带上inflight batching的整体吞吐：

如何debug

调试的话，有两个logger可以使用，也就是可以通过设置环境变量或者传入参数开启某些logger设置。

Logger could provide many useful/important information to help debugging

Python side: controlled by --log, levelin python examples (defined in tensorrt llm/logger.py)
C++ side: 这个比较隐蔽，一般是开发者使用 controlled by TLLM_LOG_LEVEL environment variable (defined in cpp/include/tensorrt llm/common/logger.h）Could print all function calls on C++ level；Help to trace the codes and locate error position

编译

这里也提到了一个加速编译的功能，有时候我们修改了一些源文件，重新编译会比较耗时。

比如改了一个.h的头文件，但是这个头文件被很多C++文件引用，所以这些个c++文件理论上都会被重编译一遍，加上trt-llm有很多kernel需要编译，编译时间很长。

官方提供了一些方法：

一般我们在某个卡测试的时候，不需要把所有cuda architecture都编译，按需编译自己当前这张卡对应architecture就行。

issue查找

可能会影响精度的选项：

用BF16训练出来，使用FP16跑（反之亦然），在小模型上可能影响不大；但是如果在大模型上，还是会有些精度问题；
context_fmha vs context_fmha_fp32_acc 默认是fp16 acc，如果遇到精度问题，可以尝试fp32_acc但是会影响速度；
Disable gemm_plugin；之前我们默认都是打开的，首先会加速编译流程；后来TRT-10优化了编译速度和支持了FP32 acc，可以尝试使用trt内部的gemm去寻找更好的性能，都可以试下；