在生产环境通过 TensorRT-LLM 部署 LLM

最新推荐文章于 2024-08-13 23:10:06 发布

程序猿李巡天

最新推荐文章于 2024-08-13 23:10:06 发布

阅读量1.2k

点赞数 18

文章标签： langchain 人工智能 AI大模型语言模型 transformer

本文链接：https://blog.csdn.net/m0_59235945/article/details/139452541

版权

引言

开源的大型语言模型确实达到了预期的效果。许多在生产环境中使用 GPT-3.5 或 GPT-4 的公司已经意识到，从成本角度来看，这些模型难以扩展。因此，企业正在寻找良好的开源替代方案。最近的模型如 Mistral 和 Llama 2 在输出质量方面表现出色。但是，将这些模型扩展以支持数千并发用户仍然是一个挑战。

尽管像 vLLM 和 TGI 这样的框架是提升推理性能的一个很好的开端，但由于缺乏一些优化，很难在生产环境中对它们进行扩展。

这就是 TensorRT-LLM 的用武之地。TensorRT-LLM 是 Nvidia 设计的一个开源框架，旨在提高大型语言模型在生产环境中的性能。像 Anthropic、OpenAI、Anyscale 等知名公司已经在使用这个框架为数百万用户服务大型语言模型。

了解 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 仓库（https://github.com/NVIDIA/TensorRT-LLM）。

使用 TensorRT-LLM 的收益

TensorRT LLM python 包允许开发者在不需要掌握 C++ 或 CUDA 的情况下，让大型语言模型达到峰值性能运行。除此之外，它还提供了令人兴奋的功能，如 token 流式传输、分页注意力和 KV 缓存。让我们深入探讨其中的一些主题。

1、分页注意力（Paged Attention）

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

在常规注意力机制中，序列的键和值必须连续存储。因此，即使你在序列内存分配的中间释放了空间，也无法将该空间用于其他序列。这会导致碎片和浪费。

标准 Attention 计算流程：所有 token 都必须保存在连续的内存块中，即使有空闲空间也无法利用。

使用分页注意力，每个页面的键/值都可以分散存储在内存中，不需连续。所以如果你在中间释放了一些页面，那个空间现在可以被其他序列重新使用。这样可以防止碎片化，提高内存利用率。在生成输出序列时，可以根据需要动态分配和释放页面。

分页注意力计算流程：之前生成的 token 可以通过删除整个页面从内存中删除。这为新的序列腾出了空间。

2、高效 KV 缓存

KV 缓存代表"键值缓存"，用于缓存大型语言模型(LLM)的部分内容，以提高推理速度并减少内存使用。

LLM 有数十亿个参数，使得在上面进行推理运行缓慢且内存密集。KV 缓存通过缓存 LLM 的层输出和激活来解决这个问题，这样就无需为每个推理重新计算。它是这样工作的:

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

理论部分已经足够了，让我们真正部署一个模型吧!

Python 动手实践

使用 TensorRT-LLM 部署模型有两个步骤：

编译模型
将编译后的模型部署为 REST API

步骤 1：编译模型

在本教程中我们将使用 Mistral 7B Instruct v0.2。如前所述，编译阶段需要 GPU。我发现在 Google Colab 笔记本上编译模型是最简单的方式。

Mistral 7B 编译器 Google Colaboratory (链接：https://colab.research.google.com/drive/1tJSMGbqYDstnChytaFb9F37oBspDcsRL)

TensorRT LLM 主要支持高端 Nvidia GPU。我在具有 A100 40GB GPU 的 Colab 上运行并将使用相同的 GPU 进行部署。

克隆 TensorRT-LLM git 仓库，其中包含编译模型所需的所有模块和脚本。

!git clone https://github.com/NVIDIA/TensorRT-LLM.git``%cd TensorRT-LLM/examples/llama

安装必要的Python依赖项。

!pip install tensorrt_llm -U --pre --extra-index-url https://pypi.nvidia.com``!pip install huggingface_hub pynvml mpi4py``!pip install -r requirements.txt

从 HuggingFace 下载 Mistral 7B Instruct v0.2 模型权重，并将它们存储在本地 tmp/hf_models/mistral-7b-instruct-v0.2 目录中。

from huggingface_hub import snapshot_download``from google.colab import userdata``   ``   ``snapshot_download(`    `"mistralai/Mistral-7B-Instruct-v0.2",`    `local_dir="tmp/hf_models/mistral-7b-instruct-v0.2",`    `max_workers=4``)

如果查看 Colab 中的 tmp/hf_models 目录，你应该可以看到模型权重。原始模型权重无法直接编译。相反，它们必须转换为特定的 TensorRT LLM 格式。

!python convert_checkpoint.py --model_dir ./tmp/hf_models/mistral-7b-instruct-v0.2 \`                             `--output_dir ./tmp/trt_engines/1-gpu/ \`                             `--dtype float16

convert_checkpoint.py 脚本将原始 Mistral 权重转换为兼容格式。
–model_dir是原始模型权重的路径。
–output_dir是转换后权重的路径。

!trtllm-build --checkpoint_dir ./tmp/trt_engines/1-gpu/ \`            `--output_dir ./tmp/trt_engines/compiled-model/ \`            `--gpt_attention_plugin float16 \`            `--gemm_plugin float16 \`            `--max_input_len 32256

trtllm-build 命令编译模型。在这个阶段，你还可以传入各种优化选项。为了简单起见，我没有使用任何额外的优化。
–checkpoint_dir是转换后模型权重的路径。
–output_dir是编译模型保存的路径。
Mistral 7B Instruct v0.2 支持 32K 的上下文长度。使用 --max_input_length 标志设置了该上下文长度。

注意：编译模型可能需要 15-30 分钟（和编译服务器硬件配置有关）。

一旦模型编译完成，你可以将编译后的模型上传到 Hugging Face Hub。为了能够上传文件到 Hugging Face Hub，你需要一个有写权限的有效访问令牌。

import os``from huggingface_hub import HfApi``   ``for root, dirs, files in os.walk(f"tmp/trt_engines/compiled-model", topdown=False):`    `for name in files:`        `filepath = os.path.join(root, name)`        `filename = "/".join(filepath.split("/")[-2:])`        `print("uploading file: ", filename)`        `api = HfApi(token=userdata.get('HF_WRITE_TOKEN'))`        `api.upload_file(`            `path_or_fileobj=filepath,`            `path_in_repo=filename,`            `repo_id="<your-repo-id>/mistral-7b-v0.2-trtllm"`        `)

将编译后的模型(.engine 文件)上传到你的 HuggingFace 账户中。
在代码中用你的 HuggingFace 用户 ID 替换

太好了！这就完成了模型编译部分。接下来是部署步骤。

步骤 2：部署编译后的模型

有许多方式可以部署这个编译后的模型。你可以使用简单的工具如 FastAPI，或更复杂的工具如 Triton 推理服务器。

当使用像 FastAPI 这样的工具时，开发人员必须自己设置 API 服务器、编写 Dockerfile 并正确配置 CUDA。管理这些可能是一件很痛苦的事情，它破坏了整体的开发体验。

为了避免这些问题，我决定使用一个简单的开源工具 Truss。Truss 允许开发人员轻松地将他们的模型与 GPU 支持打包，并在任何云环境中运行。它有许多优秀的特性，可以轻松地容器化模型：

开箱即用的 GPU 支持，无需处理 CUDA
自动创建 Dockerfile，无需自己编写
生产就绪的 API 服务器
简单的 Python 接口

使用 Truss 的主要好处是你可以轻松地将模型与 GPU 支持容器化，并将其部署到任何云环境中。

创建 Truss 环境，需要保证 Python 版本 >= 3.8

pip install --upgrade truss

可选步骤，创建 Truss 项目

truss init mistral-7b-tensort-llm

该步骤会引导你一步步完成配置，并生成项目文件。为了图省事，你可以直接使用我创建好的项目：https://github.com/htrivedi99/mistral-7b-tensorrt-llm-truss，工程目录如下：

├── mistral-7b-tensorrt-llm-truss``│   ├── config.yaml``│   ├── model``│   │   ├── __init__.py``│   │   └── model.py``|   |   └── utils.py``|   ├── requirements.txt

config.yaml 用于设置模型的各种配置,包括资源、依赖项、环境变量等。我们可以在这里指定模型名称、要安装的 Python 依赖项以及要安装的系统软件包。
model/model.py 是 Truss 的核心。它包含将在 Truss 服务器上执行的 Python 代码。在 model.py 中有两个主要方法:load()和preict()。
load 方法是我们将下载来自 hugging face 的编译模型并初始化 TensorRT LLM 引擎的地方。
predict 方法接收 HTTP 请求并调用模型。
model/utils.py 包含一些 model.py 文件的辅助函数。我没有自己编写 utils.py 文件，而是直接从 TensorRT LLM 仓库获取的。
requirements.txt 包含运行我们编译模型所需的 Python 依赖项。

深入代码解释：

model.py 包含了主要的代码，所以让我们深入探讨一下这个文件。首先看看 load 函数。

import subprocess``subprocess.run(["pip", "install", "tensorrt_llm", "-U", "--pre", "--extra-index-url", "https://pypi.nvidia.com"])``   ``import torch``from model.utils import (DEFAULT_HF_MODEL_DIRS, DEFAULT_PROMPT_TEMPLATES,`                   `load_tokenizer, read_model_name, throttle_generator)``   ``import tensorrt_llm``import tensorrt_llm.profiler``from tensorrt_llm.runtime import ModelRunnerCpp, ModelRunner``from huggingface_hub import snapshot_download``   ``STOP_WORDS_LIST = None``BAD_WORDS_LIST = None``PROMPT_TEMPLATE = None``   ``class Model:`    `def __init__(self, **kwargs):`        `self.model = None`        `self.tokenizer = None`        `self.pad_id = None`        `self.end_id = None`        `self.runtime_rank = None`        `self._data_dir = kwargs["data_dir"]``   `    `def load(self):`        `snapshot_download(`            `"htrivedi99/mistral-7b-v0.2-trtllm",`            `local_dir=self._data_dir,`            `max_workers=4,`        `)``   `        `self.runtime_rank = tensorrt_llm.mpi_rank()``   `        `model_name, model_version = read_model_name(f"{self._data_dir}/compiled-model")`        `tokenizer_dir = "mistralai/Mistral-7B-Instruct-v0.2"``   `        `self.tokenizer, self.pad_id, self.end_id = load_tokenizer(`            `tokenizer_dir=tokenizer_dir,`            `vocab_file=None,`            `model_name=model_name,`            `model_version=model_version,`            `tokenizer_type="llama",`        `)``   ``   `        `runner_cls = ModelRunner`        `runner_kwargs = dict(engine_dir=f"{self._data_dir}/compiled-model",`                             `lora_dir=None,`                             `rank=self.runtime_rank,`                             `debug_mode=False,`                             `lora_ckpt_source="hf",`                            `)``   `        `self.model = runner_cls.from_dir(**runner_kwargs)

这里发生的情况是:

在文件顶部，我们导入了必需的模块，特别是 tensorrt_llm
接下来，在 load 函数中，我们使用 snapshot_download 函数下载编译好的模型。我编译好的模型位于以下 repo id：htrivedi99/mistral-7b-v0.2-trtllm。如果你将编译好的模型上传到其他地方，请相应地更新这个值。
然后，我们使用 model/utils.py 中的 load_tokenizer 函数下载模型的 tokenizer。
最后，我们使用 TensorRT LLM 和 ModelRunner 类加载我们编译好的模型。

很酷，让我们再看一下 predict 函数。

def predict(self, request: dict):``   `        `prompt = request.pop("prompt")`        `max_new_tokens = request.pop("max_new_tokens", 2048)`        `temperature = request.pop("temperature", 0.9)`        `top_k = request.pop("top_k",1)`        `top_p = request.pop("top_p", 0)`        `streaming = request.pop("streaming", False)`        `streaming_interval = request.pop("streaming_interval", 3)``   `        `batch_input_ids = self.parse_input(tokenizer=self.tokenizer,`                                      `input_text=[prompt],`                                      `prompt_template=None,`                                      `input_file=None,`                                      `add_special_tokens=None,`                                      `max_input_length=1028,`                                      `pad_id=self.pad_id,`                                      `)`        `input_lengths = [x.size(0) for x in batch_input_ids]``   `        `outputs = self.model.generate(`            `batch_input_ids,`            `max_new_tokens=max_new_tokens,`            `max_attention_window_size=None,`            `sink_token_length=None,`            `end_id=self.end_id,`            `pad_id=self.pad_id,`            `temperature=temperature,`            `top_k=top_k,`            `top_p=top_p,`            `num_beams=1,`            `length_penalty=1,`            `repetition_penalty=1,`            `presence_penalty=0,`            `frequency_penalty=0,`            `stop_words_list=STOP_WORDS_LIST,`            `bad_words_list=BAD_WORDS_LIST,`            `lora_uids=None,`            `streaming=streaming,`            `output_sequence_lengths=True,`            `return_dict=True)``   `        `if streaming:`            `streamer = throttle_generator(outputs, streaming_interval)``   `            `def generator():`                `total_output = ""`                `for curr_outputs in streamer:`                    `if self.runtime_rank == 0:`                        `output_ids = curr_outputs['output_ids']`                        `sequence_lengths = curr_outputs['sequence_lengths']`                        `batch_size, num_beams, _ = output_ids.size()`                        `for batch_idx in range(batch_size):`                            `for beam in range(num_beams):`                                `output_begin = input_lengths[batch_idx]`                                `output_end = sequence_lengths[batch_idx][beam]`                                `outputs = output_ids[batch_idx][beam][`                                          `output_begin:output_end].tolist()`                                `output_text = self.tokenizer.decode(outputs)``   `                                `current_length = len(total_output)`                                `total_output = output_text`                                `yield total_output[current_length:]`            `return generator()`        `else:`            `if self.runtime_rank == 0:`                `output_ids = outputs['output_ids']`                `sequence_lengths = outputs['sequence_lengths']`                `batch_size, num_beams, _ = output_ids.size()`                `for batch_idx in range(batch_size):`                    `for beam in range(num_beams):`                        `output_begin = input_lengths[batch_idx]`                        `output_end = sequence_lengths[batch_idx][beam]`                        `outputs = output_ids[batch_idx][beam][`                                  `output_begin:output_end].tolist()`                        `output_text = self.tokenizer.decode(outputs)`                        `return {"output": output_text}

这里发生的是:

predict 函数接受一些模型输入，如 prompt、max_new_tokens、temperature 等。我们在函数顶部使用 request.pop 方法提取所有这些值。
接下来，我们使用 self.parse_input 辅助函数将提示格式化为 TensorRT LLM 所需的格式。
然后，我们调用我们的 LLM 模型来使用 self.model.generate 函数生成输出。generate 函数接受各种参数，帮助控制 LLM 的输出。
我还添加了一些代码来启用流式处理，通过生成一个 generator 对象。如果禁用流式处理，tokenizer 只需解码 LLM 的输出并将其作为 JSON 对象返回。

棒极了！以上就涵盖了我们关心的全部代码。下面让我们来容器化它。

容器化模型:

为了能在云端运行我们的模型，我们需要将它容器化。Truss 将负责为我们创建 Dockerfile 和打包所有内容，所以我们不需要做太多工作。

在 mistral-7b-tensorrt-llm-truss 目录之外创建一个名为 main.py 的文件。将以下代码粘贴到其中:

import truss``from pathlib import Path``   ``tr = truss.load("./mistral-7b-tensorrt-llm-truss")``command = tr.docker_build_setup(build_dir=Path("./mistral-7b-tensorrt-llm-truss"))``print(command)

运行 main.py 文件，然后查看 mistral-7b-tensorrt-llm-truss 目录。你应该会看到一堆自动生成的文件。我们不需要关心这些文件的含义,这只是 Truss 在施展它的魔力。

接下来，让我们使用 docker 构建我们的容器。按顺序运行下面的命令:

docker build mistral-7b-tensorrt-llm-truss -t mistral-7b-tensorrt-llm-truss:latest``docker tag mistral-7b-tensorrt-llm-truss <docker_user_id>/mistral-7b-tensorrt-llm-truss``docker push <docker_user_id>/mistral-7b-tensorrt-llm-truss

Sweet！我们已经准备好在云端部署模型了!

在 GKE 上部署模型

在这一节中，我们将在 Google Kubernetes Engine 上部署模型。如果还记得的话，在模型编译步骤中，我们在配备 A100 40GB GPU 的 Google Colab 上运行。对于 TensorRT LLM 来说，我们需要在完全相同的 GPU 上进行推理部署。

我不会太深入探讨如何建立 GKE 集群,因为这不在本文的范围之内。但是,我会提供我用于集群的规格。规格如下:

1 个节点，标准 kubernetes 集群(不是自动驾驶)
1.28.5 gke kubernetes 版本
1 个 Nvidia A100 40GB GPU
a2-highgpu-1g 机器(12 vCPU，85 GB 内存)
Google 管理的 GPU 驱动程序安装(否则我们需要手动安装 CUDA 驱动程序)
所有这些都将在 spot 实例上运行

配置好集群后，我们可以启动它并连接到它。集群处于活动状态且成功连接后，创建以下 kubernetes 部署:

apiVersion: apps/v1``kind: Deployment``metadata:`  `name: mistral-7b-v2-trt`  `namespace: default``spec:`  `replicas: 1`  `selector:`    `matchLabels:`      `component: mistral-7b-v2-trt-layer`  `template:`    `metadata:`      `labels:`        `component: mistral-7b-v2-trt-layer`    `spec:`      `containers:`      `- name: mistral-container`        `image: htrivedi05/mistral-7b-v0.2-trt:latest`        `ports:`          `- containerPort: 8080`        `resources:`          `limits:`            `nvidia.com/gpu: 1`      `nodeSelector:`        `cloud.google.com/gke-accelerator: nvidia-tesla-a100``---``apiVersion: v1``kind: Service``metadata:`  `name: mistral-7b-v2-trt-service`  `namespace: default``spec:`  `type: ClusterIP`  `selector:`    `component: mistral-7b-v2-trt-layer`  `ports:`  `- port: 8080`    `protocol: TCP`    `targetPort: 8080

上面是一个标准的 kubernetes 部署，它运行一个使用镜像htrivedi05/mistral-7b-v0.2-trt:latest 的容器。如果你在前一节中创建了自己的镜像，请继续使用它。否则，你可以使用我的镜像。

你可以通过运行以下命令来创建部署:

kubectl create -f mistral-deployment.yaml

为 kubernetes pod 配置资源需要几分钟时间。一旦 pod 运行，我们之前编写的 load 函数将被执行。你可以通过运行以下命令来检查 pod 的日志：

kubectl logs <pod-name>

一旦模型加载完毕，你将在pod日志中看到类似"Completed model.load() execution in 449234 ms"的消息。要通过 HTTP 向模型发送请求，我们需要端口转发服务。你可以使用以下命令进行端口转发:

kubectl port-forward svc/mistral-7b-v2-trt-service 8080

太好了！我们终于可以开始向模型发送请求了！

import requests``   ``data = {"prompt": "What is a mistral?"}``res = requests.post("http://127.0.0.1:8080/v1/models/model:predict", json=data)``res = res.json()``print(res)

你会得到如下输出：

{"output": "A Mistral is a strong, cold wind that originates in the Rhone Valley in France. It is named after the Mistral wind system, which is associated with the northern Mediterranean region. The Mistral is known for its consistency and strength, often blowing steadily for days at a time. It can reach speeds of up to 130 kilometers per hour (80 miles per hour), making it one of the strongest winds in Europe. The Mistral is also known for its clear, dry air and its role in shaping the landscape and climate of the Rhone Valley."}

LLM上传至TensorRT LLM的性能提升可以通过查看标记的流式传输来直观地感受。这里有一个流式示例:

data = {"prompt": "What is mistral wind?", "streaming": True, "streaming_interval": 3}``res = requests.post("http://127.0.0.1:8080/v1/models/model:predict", json=data, stream=True)``   ``for content in res.iter_content():`    `print(content.decode("utf-8"), end="", flush=True)

这个 Mistral 模型有相当大的上下文窗口，因此请尽情尝试使用不同的提示。

性能基准测试

仅仅看看被流式传输的标记，你可能就能感受到 TensorRT LLM 的速度有多快。不过，我还是想获得真实数字来捕捉使用 TensorRT LLM 所获得的性能提升。我运行了一些自定义基准测试，并获得了以下结果:

(提供了小/中/大提示下TensorRT LLM与Hugging Face性能对比的基准测试结果图表)

结论

在这篇博文中，我的目标是展示如何使用 TensorRT LLM 实现最先进的推理性能。我们涵盖了从编译 LLM 到在生产环境中部署模型的全过程。

虽然 TensorRT LLM 比其他推理优化器更复杂，但其性能是无可否认的。这个工具提供了 LLM 的最先进优化，同时完全开源，并针对商业使用而设计。这个框架目前还处于初期阶段，正在积极开发中。我们今天看到的性能在未来几年只会变得更好。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述