生成式人工智能/大语言模型

借助 NVIDIA TensorRT-LLM 和 NVIDIA Triton 推理服务器提升 Meta Llama 3 性能

我们很高兴地宣布支持 Meta Llama 3 系列模型,通过 NVIDIA TensorRT LLM 加速和优化 LLM 推理性能。您现在可以立即尝试 Llama 3 8B 和 Llama 3 70B,该系列中的第一个模型,通过浏览器用户界面进行体验。另外,您也可以通过 NVIDIA API 产品目录 中的 API 端点访问 Llama 3,后者被包装为 NVIDIA NIM,提供了标准 API,能够部署在任何地方。

大型语言模型需要大量的计算资源。它们的大小使其运行成本高昂且速度缓慢,尤其是在没有正确技术的情况下。许多优化技术都可用,例如内核融合量化到运行时优化(如 C++ 实现、KV 缓存、连续运行中批处理分页注意力)。 开发人员必须决定哪种组合有助于他们的用例。 TensorRT-LLM 简化了这项工作。

TensorRT-LLM 是一个开源库,用于加速 NVIDIA GPU 上最新的 LLM 推理性能。NeMo 是用于构建、自定义和部署生成式 AI 应用的端到端框架,并使用 TensorRT-LLM 和 NVIDIA Triton 推理服务器 来进行生成式 AI 部署。

TensorRT-LLM 使用 NVIDIA TensorRT 深度学习编译器。它包含最新的优化的内核,用于先进的 FlashAttention 实现,以及用于 LLM 模型执行的遮罩多头注意力 (MHA)。它还包含预处理和后处理步骤以及简单的开源 Python API 中的多 GPU/多节点通信基元,可在 GPU 上实现突破性的 LLM 推理性能。

要了解该库及其使用方法,我们来看看一个示例,了解如何使用和部署 Llama 3 8B 以及 TensorRT-LLM 和 Triton 推理服务器。

有关更深入的模型和执行细节,请参阅 TensorRT-LLM 示例

开始安装

首先,我们提供了特定于操作系统的安装步骤,以克隆和构建 TensorRT-LLM 库。您可以按照 此处说明 使用 pip 命令进行更简单的安装。此外,您也可以使用 Dockerfile 来检索依赖项并安装库。

以下命令拉取开源库并在容器内安装 TensorRT-LLM 所需的依赖项。

git clone -b v0.8.0 https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM

检索模型权重

TensorRT-LLM 是一个用于 LLM 推理的库。要使用它,您需要提供一组经过训练的权重。这些权重可以从 Hugging Face Hub 或 NVIDIA NGC 目录中获得。另一种方法是使用自己的模型权重在框架中进行训练,例如 NeMo。

本文中的命令会自动从 Hugging Face Hub 中提取 80 亿参数 Llama 3 模型的指令调整变体的权重(和分词器文件)。 您还可以使用以下命令下载要离线使用的权重,并更新后面命令中的路径以指向此目录:

git lfs install
git clone https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct

请注意,使用此模型需要遵守 特定 许可。同意条款并 使用 HuggingFace 进行身份验证 下载必要的文件。

运行 TensorRT-LLM 容器

我们将启动基础 Docker 容器,并安装 TensorRT-LLM 所需的依赖项。

# Obtain and start the basic docker image environment.
docker run --rm --runtime=nvidia --gpus all --volume ${PWD}:/TensorRT-LLM --entrypoint /bin/bash -it --workdir /TensorRT-LLM nvidia/cuda:12.1.0-devel-ubuntu22.04

# Install dependencies, TensorRT-LLM requires Python 3.10
apt-get update && apt-get -y install python3.10 python3-pip openmpi-bin libopenmpi-dev

# Install the stable version (corresponding to the cloned branch) of TensorRT-LLM.
pip3 install tensorrt_llm==0.8.0 -U --extra-index-url https://pypi.nvidia.com

编译模型

该过程的下一步是将模型编译为 TensorRT 引擎,其中包含模型权重和使用 TensorRT-LLM Python API 编写的模型定义。

TensorRT-LLM 资源库包含几个 模型架构。我们使用 Llama 模型定义。有关更多详细信息以及更可靠的插件和量化工具,请参阅此处 Llama 示例精度文档

# Log in to huggingface-cli
# You can get your token from huggingface.co/settings/token
huggingface-cli login --token *****

# Build the Llama 8B model using a single GPU and BF16.
python3 examples/llama/convert_checkpoint.py --model_dir ./Meta-Llama-3-8B-Instruct \
            --output_dir ./tllm_checkpoint_1gpu_bf16 \
            --dtype bfloat16

trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_bf16 \
            --output_dir ./tmp/llama/8B/trt_engines/bf16/1-gpu \
            --gpt_attention_plugin bfloat16 \
            --gemm_plugin bfloat16

当我们使用 TensorRT-LLM API 创建模型定义时,我们会从 TensorRT 构成神经网络层的基元。这些操作映射到 GPU 预写程序的特定内核,以实现高效的推理。

TensorRT 编译器可以横扫图形,为每个操作和每个可用 GPU 选择最佳内核。它还可以识别图形中多个操作适合合并到单个融合内核的模式,从而减少所需的内存移动量和启动多个 GPU 内核的开销。

此外,TensorRT 还将运算图构建为 NVIDIA CUDA 图形,可以同时启动。这进一步减少了启动内核的开销。

TensorRT 编译器在融合层和提高执行速度方面效率很高,但也有一些复杂的层融合,例如 FlashAttention 涉及将许多运算交织在一起且无法自动发现的操作。对于这些操作,我们可以将部分图形显式替换为 插件。在我们的示例中,我们包含 gpt_attention 插件 (可实现类似 FlashAttention 的融合注意力核函数) 和 gemm 插件 (可通过 FP32 累加执行矩阵乘法运算)。我们还将完整模型的所需精度调用为 FP16,以匹配从 HuggingFace 下载的权重的默认精度。

完成运行构建脚本后,我们有望在/tmp/llama/8B/trt_engines/bf16/1-gpu 文件夹中看到以下三个文件:

  • rank0.engine 是构建脚本的主要输出,其中包含嵌入模型权重的可执行操作图。
  • `config.json` 文件包含模型的详细信息,例如其结构和精度,以及引擎中集成的插件列表。

运行模型

既然我们已经有了模型引擎,我们可以用它做什么?

引擎文件包含执行模型所需的信息。TensorRT-LLM 包含高度优化的 C++运行时,用于执行引擎文件和管理流程,例如从模型输出中采样令牌、管理 KV 缓存以及同时处理批处理请求。

我们可以直接使用运行时环境在本地执行模型,也可以在生产环境中使用 Triton 推理服务器进行部署,以便与多个用户共享模型。

要在本地运行模型,我们可以执行以下命令:

python3 examples/run.py --engine_dir=./tmp/llama/8B/trt_engines/bf16/1-gpu --max_output_len 100 --tokenizer_dir ./Meta-Llama-3-8B-Instruct --input_text "How do I count to nine in French?"

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

除了本地执行,我们还可以使用 Triton 推理服务器 以创建 LLM 的生产就绪型部署。适用于 TensorRT-LLM 的 Triton 推理服务器后端 使用 TensorRT-LLM C++运行时实现高性能推理执行。它包括 动态批处理分页 KV 缓存,可提供低延迟和高吞吐量。TensorRT-LLM 后端已与 Triton 推理服务器捆绑在一起,并可作为预构建容器使用 在 NGC 上运行。

首先,我们必须创建模型库,以便 Triton 推理服务器能够读取模型和任何相关元数据。

tensorrtllm_backend 存储库包含所需的模型库,位于 `all_models/inflight_batcher_llm/` 中。

该目录包含四个子文件夹,分别包含模型执行过程的不同部分的构件。Triton 推理服务器 Python 后端 中的 `preprocessing/` 和 `postprocessing/` 文件夹包含用于对文本输入进行标记化和去标记化的脚本,以在字符串和模型所运行的令牌 ID 之间进行转换。

我们的 `tensorrt_llm` 文件夹包含之前编译的模型引擎。最后,我们的 `ensemble` 文件夹定义了 模型集成,该技术将之前的三个组件连接在一起,并指示 Triton 推理服务器如何通过这些组件流数据。

向下拉取示例模型库,并将您在上一步中编译的模型复制到其中。

# After exiting the TensorRT-LLM docker container
cd ..
git clone -b v0.8.0 https://github.com/triton-inference-server/tensorrtllm_backend.git
cd tensorrtllm_backend
cp ../TensorRT-LLM/tmp/llama/8B/trt_engines/bf16/1-gpu/* all_models/inflight_batcher_llm/tensorrt_llm/1/

接下来,我们必须使用编译模型引擎的位置修改资源库骨架中的配置文件。我们还必须更新配置参数 (例如 tokenizer),以便在批处理推理请求时使用和处理 KV 缓存的内存分配。

#Set the tokenizer_dir and engine_dir paths
HF_LLAMA_MODEL=TensorRT-LLM/Meta-Llama-3-8B-Instruct
ENGINE_PATH=tensorrtllm_backend/all_models/inflight_batcher_llm/tensorrt_llm/1

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/preprocessing/config.pbtxt tokenizer_dir:${HF_LLAMA_MODEL},tokenizer_type:auto,triton_max_batch_size:64,preprocessing_instance_count:1

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/postprocessing/config.pbtxt tokenizer_dir:${HF_LLAMA_MODEL},tokenizer_type:auto,triton_max_batch_size:64,postprocessing_instance_count:1

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/tensorrt_llm_bls/config.pbtxt triton_max_batch_size:64,decoupled_mode:False,bls_instance_count:1,accumulate_tokens:False

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/ensemble/config.pbtxt triton_max_batch_size:64

python3 tools/fill_template.py -i all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt triton_max_batch_size:64,decoupled_mode:False,max_beam_width:1,engine_dir:${ENGINE_PATH},max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:True,enable_kv_cache_reuse:False,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0

现在,我们可以启动 Docker 容器并启动 Triton 服务器。我们必须指定 **世界大小** – 模型构建所使用的 GPU 数量 – 并指向我们刚刚设置的 model_repo。

#Change to base working directory
cd..
docker run -it --rm --gpus all --network host --shm-size=1g \
-v $(pwd):/workspace \
--workdir /workspace \
nvcr.io/nvidia/tritonserver:24.03-trtllm-python-py3

# Log in to huggingface-cli to get tokenizer
huggingface-cli login --token *****

# Install python dependencies
pip install sentencepiece protobuf

# Launch Server

python3 tensorrtllm_backend/scripts/launch_triton_server.py --model_repo tensorrtllm_backend/all_models/inflight_batcher_llm --world_size 1

发送请求

要发送推理请求并从正在运行的服务器接收完成结果,您可以使用 Triton 推理服务器客户端库 或将 HTTP 请求发送到 生成的端点。以下 curl 命令演示了正在运行的服务器请求完成等操作的快速测试:功能齐全的客户端脚本 可以进行审查,以便与服务器进行通信。

curl -X POST localhost:8000/v2/models/ensemble/generate -d \
'{
"text_input": "How do I count to nine in French?",
"parameters": {
"max_tokens": 100,
"bad_words":[""],
"stop_words":[""]
}
}'

结束语

TensorRT-LLM 提供用于在 NVIDIA GPU 上优化和高效运行大型语言模型的工具。Triton 推理服务器是部署和高效服务 Llama 3 等大型语言模型的理想选择。

使用此入门指南,并利用开源工具开始您的使用之旅,以使用 Llama 3 和许多其他大型语言模型。

NVIDIA AI Enterprise 包含 TensorRT 的端到端 AI 软件平台 TensorRT,即将包含 TensorRT-LLM,以提供具有企业级安全性、稳定性、可管理性和支持的任务关键型 AI 推理。

入门资源

  • 访问 TensorRT-LLM 开源库。
  • 详细了解 NVIDIA NeMo 开源库的 功能和特性
  • 阅读 TensorRT 和 TensorRT-LLM 开发者指南。
  • 在 GitHub 上探索我们的示例代码、基准测试和文档
  • NVIDIA NIM 利用 TensorRT-LLM 优化推理,以在 ai.nvidia.com 上提供高效的模型推理。

Tags