NVIDIA TensorRT LLM 优化大型语言模型推理，现已公开

这篇文章最初发表在 NVIDIA 技术博客上。

NVIDIA 于 2023 年 10 月 19 日公开发布 TensorRT-LLM，可在 NVIDIA GPU 上加速和优化最新的大语言模型（Large Language Models）的推理性能。该开源程序库现已作为 NVIDIA NeMo 框架的一部分，在 /NVIDIA/TensorRT-LLM GitHub 资源库中免费提供。

大语言模型彻底改变了人工智能领域，并创造了与数字世界交互的全新方式。但是，随着世界各地的企业和应用开发者想要在自己的工作中采用大语言模型，运行这些模型的相关问题开始凸显。

简而言之，大语言模型规模庞大。如果不采用正确的技术，那么运行大语言模型的成本不仅会很高，速度也会很慢。

为解决这一问题，从模型优化（如内核融合和量化）到运行时优化（如 C++ 实现、KV 缓存、连续动态批处理 continuous in-flight batching 和分页注意力 paged attention），众多优化技术应运而生。但很难确定哪种技术适合您的用例，也很难在这些实施中有时并不兼容的技术间进行交互。

因此，NVIDIA 推出了 TensorRT-LLM，它是一个专门用于编译和优化大语言模型推理的综合程序库。TensorRT-LLM 整合了所有这些优化功能，同时还提供了一个直观的 Python API 来定义和构建新模型。

TensorRT-LLM 开源程序库可加快 NVIDIA GPU 上最新大语言模型的推理性能。它是 NVIDIA NeMo 中优化大语言模型推理的骨干力量。NeMo 是一个用于构建和定制生成式 AI 应用并将其部署到生产中的端到端框架，为生成式 AI 的部署提供了完整的容器，如 TensorRT-LLM 和 NVIDIA Triton 等。

TensorRT-LLM 现在也可作为测试版本用于原生 Windows 系统。应用开发者和 AI 爱好者现可在本地运行由 NVIDIA RTX 和 NVIDIA GeForce RTX GPU 驱动的个人电脑和工作站上加速大语言模型。

TensorRT-LLM 包含 TensorRT 的深度学习编译器，并采用了最新的优化内核，这些内核专为在前沿研究中实施用于 大语言模型 执行的 FlashAttention 和带掩码的多头注意力 (masked multi-head attention) 而打造。

TensorRT-LLM 还将预处理和后处理步骤以及多 GPU /多节点通信基元整合到一个简单的开源 Python API 中，可在 GPU 上实现突破性的大语言模型推理性能。

TensorRT-LLM 的特性如下：

• 支持 Llama 1 和 2、Bloom、ChatGLM、Falcon、MPT、Baichuan 及 Starcoder 等 大语言模型
• 动态批处理和分页注意力
• 多 GPU 多节点（MGMN）推理
• FP8 精度的 NVIDIA Hopper Transformer 引擎
• 支持 NVIDIA Ampere 架构、NVIDIA Ada Lovelace 架构和 NVIDIA Hopper GPU
• 原生 Windows 支持（测试版）

过去两年，NVIDIA 一直与领先的大语言模型企业与团队开展密切合作，以加速和优化大语言模型 推理，包括 Anyscale、百度搜索、百川智能、BOSS直聘 、Cohere、Deci、Grammarly、吉利汽车研究院、金山办公、京东探索研究院、蚂蚁集团、美团、Meta、Mistral AI、MosaicML（现已并入 Databricks）、OctoML、Perplexity AI、PPIO边缘云、Tabnine、Together.ai、Unisound、微博、蔚来汽车、小鹏汽车语音开发团队、智谱 AI 、智源、中科创达等（公司名按首字母英文排序）。

为帮助您了解该程序库及其使用方式，可通过一个示例来了解如何在 Linux 上通过 TensorRT-LLM 和 NVIDIA Triton 使用并部署 Llama 2（一个常用且公开的大语言模型）。如要使用 Windows 测试版，参见用于原生 Windows 的 TensorRT-LLM GitHub 资源库。

如需了解不同模型、不同优化以及多 GPU 执行等更多信息，请参见完整的 TensorRT-LLM 示例列表。

开始安装

首先克隆并构建 TensorRT-LLM 程序库。构建 TensorRT-LLM 并检索其所有依赖项的最简单方法是使用随附的 Dockerfile：

git clone -``b release``/``0.5``.``0 https:``/``/``github.com``/``NVIDIA``/``TensorRT``-``LLM.git
cd TensorRT``-``LLM
git submodule update -``-``init -``-``recursive
make -``C docker release_build

这些命令会提取一个基础容器，并在容器内安装 TensorRT-LLM 所需的所有依赖项，然后在容器中构建并安装 TensorRT-LLM。

检索模型权重

TensorRT-LLM 是一个用于大语言模型推理的程序库，因此要使用它，就需要提供一组训练过的权重。您可以使用自己在 NVIDIA NeMo 等框架中训练的模型权重，也可以从 HuggingFace Hub 等资源库中提取一组预训练权重。

本文中的命令会自动从 HuggingFace Hub 中提取 70 亿参数的 Llama 2 模型聊天调优变体的权重和分词器文件。您还可以使用以下命令，自行下载权重以供离线使用。您只需更新后续命令中的路径，使其指向此目录即可：

git lfs install
git clone https:``/``/``huggingface.co``/``meta``-``llama``/``Llama``-``2``-``7b``-``chat``-``hf

该模型的使用受特定许可的限制。

编译模型

下一步是将模型编译到 TensorRT 引擎中。为此，像定义模型权重那样，您也需要在 TensorRT-LLM Python API 中编写模型定义。

TensorRT-LLM 资源库包含多种预定义模型架构。在本文中，您可以使用其中包含的 Llama 模型定义，而无需自行编写。下面是 TensorRT-LLM 中一些可用优化的最简示例。

有关可用插件和量化的更多信息，参见完整的 Llama 示例和数值精度。

# Launch the Tensorrt-LLM container
make -``C docker release_run LOCAL_USER``=``1
# Log in to huggingface-cli
# You can get your token from huggingface.co/settings/token
huggingface``-``cli login -``-``token *``*``*``*``*
# Compile model
python3 examples``/``llama``/``build.py \
-``-``model_dir meta``-``llama``/``Llama``-``2``-``7b``-``chat``-``hf \
-``-``dtype float16 \
-``-``use_gpt_attention_plugin float16 \
-``-``use_gemm_plugin float16 \
-``-``remove_input_padding \
-``-``use_inflight_batching \
-``-``paged_kv_cache \
-``-``output_dir examples``/``llama``/``out

使用 TensorRT-LLM API 创建模型定义时，可以使用构成神经网络层的 NVIDIA TensorRT 基元来构建操作图。这些操作会映射到特定的内核，即为 GPU 预写的程序。

TensorRT 编译器可以对图进行扫描，为每个操作和可用的 GPU 选择最佳内核。最重要的是，它还能在图中识别出可以将多个操作*"融合"*到一个内核中的模式。这就减少了启动多个 GPU 内核所需的显存移动量和开销。

TensorRT 还能将操作图编译成单个 CUDA Graph，其可以一次性启动，从而进一步减少内核启动开销。

TensorRT 编译器在融合多个层和提高执行速度方面非常强大，但有些复杂的层融合（如 FlashAttention 等）会将许多操作交叉在一起，而且无法被自动发现。对此，可以在编译时使用插件来对图中的部分内容进行显式替换。

在本例中，您加入了 gpt_attention 插件（实现类似 FlashAttention 的融合注意力内核）和 gemm 插件（使用 FP32 累加执行矩阵乘法）。您还可以将完整模型的期望精度设为 FP16，以便与您从 HuggingFace 下载的权重的默认精度相匹配。

下面是该脚本运行后所生成的结果。现在，在 /examples/llama/out 文件夹中有以下文件：

• Llama_float16_tp1_rank0.engine：构建脚本的主要输出，包含嵌入模型权重的可执行操作图。
• config.json：包含模型的详细信息，如模型的一般结构和精度以及关于引擎中包含哪些插件的信息。
• model.cache：缓存模型编译过程中的部分时序和优化信息，使后续构建更快。

运行模型

现在您已经有了模型引擎，接下来该用它做什么呢？

引擎文件包含执行模型所需的信息，但在实际使用大语言模型时，需要的远不止是一次通过模型的前向传播。TensorRT-LLM 包含一个高度优化的 C++ 运行时，以执行已构建的大语言模型引擎和管理若干流程，如从模型输出结果中采样分词、管理 KV 缓存和批量处理请求等。

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

如要在本地运行模型，请执行以下命令：

python3 examples``/``llama``/``run.py -``-``engine_dir``=``examples``/``llama``/``out -``-``max_output_len 100 -``-``tokenizer_dir meta``-``llama``/``Llama``-``2``-``7b``-``chat``-``hf -``-``input_text "How do I count to nine in French?"

使用 Triton 理服务器部署

除了在本地执行之外，您还可以使用 NVIDIA Triton 推理服务器为您的大语言模型创建生产就绪部署。

NVIDIA为 TensorRT-LLM 发布了一个新的 Triton 推理服务器后端。该后端利用 TensorRT-LLM C++ 运行时实现快速推理执行，并加入动态批处理和分页 KV 缓存等技术。

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

cd ..
git clone -``b release``/``0.5``.``0
git@github.com:triton``-``inference``-``server``/``tensorrtllm_backend.git
cd tensorrtllm_backend
git submodule update -``-``init -``-``recursive
git lfs install
git lfs pull
DOCKER_BUILDKIT``=``1 docker build -``t triton_trt_llm -``f dockerfile``/``Dockerfile.trt_llm_backend .

接下来，创建一个模型资源库，以便 Triton 推理服务器能够读取模型和任何相关元数据。tensorrtllm_backend 资源库包含合适模型资源库框架，该框架位于 all_models/inflight_batcher_llm/ 目录下以供使用。现在该目录下有四个子文件夹，分别存放模型执行流程中各个部分的工件：

• /preprocessing 和 /postprocessing：包含 Triton 推理服务器 Python 后端的脚本，用于将文本输入分词话，并对模型输出结果进行反分词化处理，以便在字符串和模型使用的分词 ID 之间进行转换。
• /tensorrt_llm：用于放置之前编译的模型引擎。
• /ensemble：定义一个模型集合，将前面三个组件连接在一起并告诉 Triton 推理服务器如何通过其传输数据。

将编译好的模型复制到模型资源库：

rm all_models``/``inflight_batcher_llm``/``tensorrt_llm``/``1``/``*
cp tensorrt_llm``/``examples``/``llama``/``out``/``* all_models``/``inflight_batcher_llm``/``tensorrt_llm``/``1``/

现在使用以下信息修改资源库框架中的一些配置文件：

• 编译后的模型引擎的位置
• 待使用的分词器
• 成批执行推理时如何处理 KV 缓存的内存分配

python3 tools``/``fill_template.py -``-``in_place \
all_models``/``inflight_batcher_llm``/``tensorrt_llm``/``config.pbtxt \
decoupled_mode:true,engine_dir:``/``app``/``all_models``/``inflight_batcher_llm``/``tensorrt_llm``/``1``,\
max_tokens_in_paged_kv_cache:,batch_scheduler_policy:guaranteed_completion,kv_cache_free_gpu_mem_fraction:``0.2``,\
max_num_sequences:``4
python tools``/``fill_template.py -``-``in_place \
all_models``/``inflight_batcher_llm``/``preprocessing``/``config.pbtxt \
tokenizer_type:llama,tokenizer_dir:meta``-``llama``/``Llama``-``2``-``7b``-``chat``-``hf
python tools``/``fill_template.py -``-``in_place \
all_models``/``inflight_batcher_llm``/``postprocessing``/``config.pbtxt \
tokenizer_type:llama,tokenizer_dir:meta``-``llama``/``Llama``-``2``-``7b``-``chat``-``hf

最后，您可以启动 Docker 容器和 Triton 服务器，指定要部署在多少个 GPU 上（总进程数量），并指向刚刚设置的 model_repo 值。

docker run -``it -``-``rm -``-``gpus all -``-``network host -``-``shm``-``size``=``1g -``v $(pwd)``/``all_models:``/``app``/``all_models triton_trt_llm
# Log in to huggingface-cli to get the tokenizer
huggingface``-``cli login -``-``token *``*``*``*``*
# Launch the server
python3 scripts``/``launch_triton_server.py -``-``world_size 1 -``-``model_repo /``app``/``all_models``/``inflight_batcher_llm

发送请求

如要向运行中的服务器发送请求并与之交互，您可以使用一个 Triton 推理服务器客户端程序库，这些程序库适用于多种不同的编程语言和环境。在初始环节，请使用下面的 Python 代码示例或 GitHub 上功能更全面的端到端客户端流式传输。

import numpy as np
import tritonclient.http as httpclient
from tritonclient.utils import np_to_triton_dtype
def prepare_tensor(name, input``):
t = httpclient.InferInput(name, input``.shape,
np_to_triton_dtype(``input``.dtype))
t.set_data_from_numpy(``input``)
return t
TRTION_URL = "localhost:8000"
client = httpclient.InferenceServerClient(TRITON_URL)
prompt = "How do I count to nine in French?"
inputs = [
prepare_tensor(``"text_input"``, np.array([[prompt]], dtype``=``object``)),
prepare_tensor(``"max_tokens"``, np.array([[``100``]], dtype``=``np.uint32)),
prepare_tensor(``"bad_words"``, np.array([[""]], dtype``=``object``)),
prepare_tensor(``"stop_words"``, np.array([[""]], dtype``=``object``))
]
result = client.infer(``"ensemble"``, inputs)
print``(result)

总结

TensorRT-LLM 和 Triton 推理服务器是优化、部署和高效运行大语言模型不可或缺的工具。随着 TensorRT-LLM 以开源程序库的形式发布在 GitHub 上，企业和应用开发者能够比以往更容易地挖掘这些模型的潜力。

如果您想要深入了解 LLM 的世界，现在正是开始使用 TensorRT-LLM 的好时机。您可以探索它的功能，尝试不同的模型和优化方法，并开启您的旅程，解锁 AI 驱动的语言模型所蕴含的惊人力量。

更多有关 TensorRT-LLM 的入门信息，参见以下资源：

• 访问 /NVIDIA/TensorRT-LLM GitHub 资源库中的开源程序库。
• 进一步了解为生成式 AI 的部署提供完整容器（包括 TensorRT-LLM 和 NVIDIA Triton）的 NVIDIA NeMo。
• 在 GitHub 上探索示例代码、基准和 TensorRT-LLM 文档。
• 购买 NVIDIA AI Enterprise。该端到端 AI 软件平台包含 TensorRT，并将在不久后加入 TensorRT-LLM，其通过提供企业级的安全性、稳定性、可管理性和支持来满足任务关键型 AI 推理的需求。欲了解更多信息，请联系销售人员。

探索 TensorRT 和 TensorRT-LLM 的入门资源。

阅读原文