TensorRT-LLM初探(三)最佳部署实践

科技   2024-08-18 23:35   日本  
 


本篇文章根据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 功能与架构

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比较被诟病的就是开源开的不彻底,能改动的地方不多。

trt-llm之开源与半开源

究其原因,是有很多针对不同硬件做了定制化优化的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接口:

image

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去寻找更好的性能,都可以试下;

如何添加一个新的模型

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]
参考资料
[1]

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

往期回顾


GiantPandaCV
专注于大语言模型,CUDA,编译器,工程部署和优化等多个方向技术分享。我们不仅坚持原创,也规范转载知乎大佬们的高质量博文。希望在传播知识、分享知识的同时能够启发你,在人类通往AGI的道路上互相帮助(・ω\x26lt;)☆
 最新文章