英伟达开源TensorRT-LLM，可优化类ChatGPT开源模型！

专注AIGC领域的专业社区，关注OpenAI、百度文心一言等大语言模型（LLM）的发展和应用落地，关注LLM的基准评测和市场研究，欢迎关注！

10月20日，全球AI领导者英伟达（NVIDIA）在官网开源了TensorRT-LLM，以加速和优化NVIDIA GPU上最新大语言模型的推理性能。

目前，优化大语言模型的方法有很多，例如，内核融合和量化、知识蒸馏、KV缓存、模型剪枝、连续运行中批处理和分页注意力等，但开发者面对如此多的技术很难选择适合自己的用例，同时无法实现更好的兼容，而TensorRT-LLM提供了统一的解决方案。

TensorRT-LLM支持Llama 1/ 2、Baichuan（百川智能）、ChatGLM、Falcon、MPT、和Starcoder等市面上高性能类ChatGPT开源模型。

开源地址：https://github.com/NVIDIA/TensorRT-LLM/tree/release/0.5.0

TensorRT-LLM简单介绍

TensorRT-LLM是一个用于编译和优化大语言模型推理的综合库。TensorRT-LLM融合了目前主流优化方法，同时提供了用于定义和构建新模型的直观Python API。

TensorRT-LLM封装了TensorRT的深度学习编译器，并包含最新的优化内核，用于实现FlashAttention和用于大语言模型执行的屏蔽多头注意力 (MHA)。

此外，TensorRT-LLM还原生支持Windows，可以帮助开发者在英伟RT和GeForce RTX GPU的PC和工作站上更快、更高效地运行大语言模型。

TensorRT-LLM主要技术亮点如下：

• 支持Llama 1/ 2、ChatGLM、Falcon、MPT、Baichuan (百川智能)和Starcoder等主流开源大语言模型。
• 连续批处理和分页注意力
• 多GPU多节点(MGMN)推理
• 支持霍普架构的transformer引擎与FP8
• 支持安培架构、阿达Lovelace架构和霍普GPU
• 原生支持Windows (测试版)

TensorRT-LLM使用示例

下面展示了TensorRT-LLM如何在Linux上使用、部署和优化Llama 2。这里需要用到NVIDIA Triton（推理服务器）。

1）开始安装

首先克隆并构建TensorRT-LLM库。构建TensorRT-LLM并检索所有依赖项的最简单方法是使用包含Dockerfile。并在容器内安装TensorRT-LLM 所需的所有依赖项。

2）检索模型权重

要使用TensorRT-LLM，需要提供一组经过训练的权重。可以使用在NVIDIA NeMo等框架中训练的自己的模型权重，也可以从 HuggingFace Hub 等存储库中提取一组预训练的权重。

我们从HuggingFaceHub中提取 7B 参数 Llama 2 模型的聊天调整变体的权重和分词器文件。

3）编译模型

需要将模型编译为TensorRT引擎，所以，要将模型权重以及用 TensorRT-LLM Python API 编写的模型进行定义。

而TensorRT-LLM 存储库包含各种预定义的模型架构，已经附带了Llama 2的模型定义，所以，开发者无需自己编写模型定义了。这也突出了TensorRT-LLM的优化功能。

此外，TensorRT-LLM编译器可以扫描图表，为每个操作和可用 GPU 选择最佳内核。还可以识别图中的模式，其中多个操作是融合到单个内核中的良好候选者。

这极大减少了所需的内存移动量以及启动多个GPU内核的消耗，对于优化大语言模型推理至关重要。

4）运行模型

开发者可以直接使用本地模型，也可以使用 NVIDIA Triton 推理服务器的 TensorRT-LLM 运行时后端为多个用户提供的模型。

5）使用 NVIDIA Triton 推理服务器进行部署

除了本地执行之外，开发者还可以使用NVIDIA Triton 推理服务器来部署大语言模型的生产环境。

首先，为Triton推理服务器构建 TensorRT-LLM后端，使用Dockerfile在容器中构建后端。

接下来，创建一个模型存储库，以便 Triton Inference Server 可以读取模型和任何关联的元数据。将编译的模型复制到模型存储库。

使用下面信息修改存储库框架中的一些配置文件：

• 编译好的模型引擎在哪里
• 使用什么标记器
• 批量推理时如何处理KV缓存的内存分配

最后，启动 Docker容器并启动 Triton 服务器，指定要部署的 GPU 数量（全局大小）并指向model_repo刚刚设置的值。

总体来说，TensorRT-LLM在大语言模型的模型编译、资源消耗、流程简化等方面表现非常优秀，同时可以帮助开发人员快速尝试不同类型的开源类ChatGPT模型。

本文素材来源英伟达官网，如有侵权请联系删除

END