本篇文章根据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, level in 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去寻找更好的性能,都可以试下;
如何添加一个新的模型
TRT-LLM使用新的Python-API的初衷就是想要后续改动或者添加新模型更方便些。
因为TRT-LLM来源于TRT,因此构建网络想要通过trt-python-api去构建,这里trt-llm对这个api做了改进,但是相比pytorch可能还是难度大些。不过trt-llm已经提供了一些例子,比如我们可以使用llama的实现去适配其他模型。
以下是官方提供的添加新模型的流程:
具体来说,就是仿照llama的实现,以及一些llm的基本class和内部具体实现的layer:
除了模型的搭建,还需要实现convert权重相关的地方,从huggingface权重到trt-llm权重格式的转换:
如果官方提供的例子没有模型中某些层的实现,但你这个层可以通过官方已经提供的layer接口实现,那么我们可以利用官方提供的Python-API搭建出来:
当然,如果提供的layer、functional接口也没有你的实现,那就只能自己搓一个kernel出来,不过这个会比较复杂。参考之前TRT中的方式,我们需要先定义一个kernel plugin,写好相关的kernel实现,然后在外部引用:
以上就是NVIDIA-AI技术开放日关于TRT-LLM 最佳性能实践的全部内容。
参考
https://www.bilibili.com/video/BV1aT42167mk/?spm_id_from=pageDriver&vd_source=eec038509607175d58cdfe2e824e8ba2[7]
NVIDIA AI技术开放日 2024 夏: https://space.bilibili.com/1320140761/channel/collectiondetail?sid=3446369
[2]TRT-LLM 最佳部署实践: https://www.bilibili.com/video/BV1MS421d7Jm/
[3]TRT-LLM 最佳部署实践: https://www.bilibili.com/video/BV1MS421d7Jm/
[4]TensorRT-LLM初探(一)基于最新commit运行llama,以及triton-tensorrt-llm-backend: https://ai.oldpan.me/t/topic/260
[5]TensorRT-LLM初探(二)简析了结构,用的更明白: https://ai.oldpan.me/t/topic/203
[6]sglang: sglang
[7]https://www.bilibili.com/video/BV1aT42167mk/?spm_id_from=pageDriver&vd_source=eec038509607175d58cdfe2e824e8ba2: https://www.bilibili.com/video/BV1aT42167mk/?spm_id_from=pageDriver&vd_source=eec038509607175d58cdfe2e824e8ba2