输入“/”快速插入内容

使用TensorRT-LLM进行生产环境的部署指南

2024年7月19日修改

作者：P**nHub兄弟网站 DeepHub IMBA

TensorRT-LLM是一个由Nvidia设计的开源框架，用于在生产环境中提高大型语言模型的性能。该框架是基于 TensorRT 深度学习编译框架来构建、编译并执行计算图，并借鉴了许多 FastTransformer 中高效的 Kernels 实现，并且可以利用 NCCL 完成设备之间的通讯。​

虽然像vLLM和TGI这样的框架是增强推理的一个很好的起点，但它们缺乏一些优化，因此很难在生产中扩展它们。所以Nvidia在TensorRT的基础上有开发了TensorRT-LLM，像Anthropic, OpenAI, Anyscale等大公司已经在使用这个框架为数百万用户提供LLM服务。​

TensorRT-LLM

与其他推理技术不同，TensorRT LLM不使用原始权重为模型服务。它会编译模型并优化内核，这样可以在Nvidia GPU上有效地服务。运行编译模型的性能优势远远大于运行原始模型。这是TensorRT LLM非常快的主要原因之一。​

原始模型权重和优化选项(如量化级别、张量的并行性、管道并行性等)一起传递给编译器。然后编译器获取该信息并输出针对特定GPU优化的模型二进制文件。​

但是这里整个模型编译过程必须在GPU上进行。生成的编译模型也是专门针对运行它的GPU进行优化的。例如，在A40 GPU上编译模型，则可能无法在A100 GPU上运行它。所以无论在编译过程中使用哪种GPU，都必须使用相同的GPU进行推理。​

但是TensorRT LLM并不支持开箱即用所有的大型语言模型（原因是每个模型架构是不同的）。但是TensorRT所作的做深度图级优化是支持大多数流行的模型，如Mistral、Llama和Qwen等。具体支持的模型可以参考TensorRT LLM Github官方的列表​

TensorRT-LLM的好处

TensorRT LLM python包允许开发人员在不了解c++或CUDA的情况下以最高性能运行LLM。​

分页注意力

大型语言模型需要大量内存来存储每个令牌的键和值。随着输入序列变长，这种内存使用会变得非常大。​

通常情况下，序列的键和值必须连续存储。所以即使你在序列的内存分配中释放了空间，你也不能把这个空间用于其他序列。这会导致碎片化和浪费。​

分页注意力将键/值分成而不是连续的页，这样可以放在内存中的任何地方，如果您在中间释放一些分页，那么这些空间可以用于其他序列。​

这可以防止碎片，并允许更高的内存利用率。在生成输出序列时，可以根据需要动态地分配和释放页面。​

高效KV缓存

llm有数十亿个参数，这使得它们运行推理时速度缓慢且占用大量内存。KV缓存通过缓存LLM的层输出和激活来帮助解决这个问题，因此它们不需要为每个推理重新计算。​

下面是它的工作原理:

在推理期间，当LLM执行每一层时，输出将被缓存到具有唯一键的键值存储中。当后续推断使用相同的层输入时，不是重新计算层，而是使用键检索缓存的输出。这避免了冗余计算，减少了激活内存，提高了推理速度和内存效率。​

下面我们开始使用TensorRT-LLM部署一个模型

TensorRT-LLM部署教程

使用TensorRT-LLM部署模型首先就是要对模型进行编译，这里我们将使用Mistral 7B instruction v0.2。编译阶段需要GPU，所以为了方便使用我们直接在Colab上操作。​

使用TensorRT-LLM进行生产环境的部署指南​