介绍

开源大型语言模型果然名不虚传。许多在生产中使用 GPT-3.5 或 GPT-4 的公司已经意识到，从成本角度来看，这些模型根本无法扩展。因此，企业正在寻找优秀的开源替代方案。最近，Mixtral 和 Llama 2 等模型在输出质量方面取得了令人瞩目的成果。但是，将这些模型扩展到支持数千名并发用户仍然是一个挑战。

虽然 vLLM 和 TGI 等框架是增强推理的良好起点，但它们缺乏一些优化功能，因此很难在生产中进行扩展。

这就是 TensorRT-LLM 的用武之地。TensorRT-LLM 是由 Nvidia 设计的开源框架，用于在生产环境中提升大型语言模型的性能。Anthropic、OpenAI、Anyscale 等大多数大公司已经在使用该框架为数百万用户提供 LLM 服务。

了解 TensorRT-LLM

与其他推理技术不同，TensorRT LLM 不使用原始权重为模型提供服务。相反，它编译模型并优化内核，以便在 Nvidia GPU 上高效运行。运行编译模型的性能优势远远大于运行原始模型。这也是 TensorRT LLM 快速运行的主要原因之一。

原始模型权重以及量化级别、张量并行性、流水线并行性等优化选项会被传递给编译器。然后，编译器根据这些信息输出针对特定 GPU 优化的模型二进制文件。

需要注意的是，整个模型编译过程必须在 GPU 上进行。生成的编译模型将专门针对其运行的 GPU 进行优化。例如，如果在 A40 GPU 上编译模型，就无法在 A100 GPU 上运行。因此，无论编译时使用的是哪种 GPU，推理时都必须使用相同的 GPU。

TensorRT LLM 并不支持所有开箱即用的大型语言模型。原因是每个模型的架构各不相同，而 TensorRT 会进行深度图级优化。尽管如此，TensorRT 还是支持大多数流行的模型，如 Mistral、Llama 和 Qwen。

使用 TensorRT-LLM 的好处

TensorRT LLM python 软件包允许开发人员以最高性能运行 LLM，而无需了解 C++ 或 CUDA。此外，它还具有令牌流、分页关注和 KV 缓存等便利功能。下面让我们深入探讨其中的几个主题。

分页关注

大型语言模型需要大量内存来存储每个标记的键和值。随着输入序列越来越长，内存使用量也会越来越大。

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

使用分页注意时，每一页键/值都可以放置在内存中的任何位置，不会相互连接。因此，如果中间腾出一些页面，这些空间就可以重新用于其他序列。

这样可以防止碎片化，提高内存利用率。在生成输出序列时，可以根据需要动态分配和释放页面。

高效 KV 缓存

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

LLM 有数十亿个参数，因此在其上运行推理既慢又耗费内存。KV 缓存通过缓存 LLM 的层输出和激活来帮助解决这个问题，这样每次推理时就不需要重新计算这些参数。

Python 实践教程

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

1. 编译模型
2. 将编译后的模型部署为 REST API 端点

第一步：编译模型

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

TensorRT LLM 主要支持高端 Nvidia GPU。我在 A100 40GB GPU 上运行了谷歌合作项目，并将使用相同的 GPU 进行部署。

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

• 克隆 TensorRT-LLM git 仓库。该 repo 包含编译模型所需的所有模块和脚本。

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

• 安装必要的 Python 依赖项。

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
)

• 从 hugging face 下载 Mistral 7B Instruct v0.2 模型权重，并将其存储在本地目录 tmp/hf_models/mistral-7b-instruct-v0.2 中。
• 在 Colab 的 tmp/hf_models 目录中查看模型权重。

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

• 原始模型权重无法编译。相反，它们必须转换成特定的 tensorRT LLM 格式。
• 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 标志设置了上下文长度。

编译好模型后，就可以将编译好的模型上传到 hugging face hub。要将文件上传到 hugging face hub，你需要一个具有 WRITE 访问权限的有效访问令牌。

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 文件）上传到 hugging face，用户名是你的用户 ID。
• 将代码中的 <your-repo-id> 替换为你的 hugging face repo，通常也就是你的 hugging face 用户 ID。

第二步：部署编译后的模型

有很多方法可以部署编译后的模型。你可以使用像 FastAPI 这样的简单工具，也可以使用像 triton 推论服务器这样的复杂工具。

在使用 FastAPI 等工具时，开发人员必须设置 API 服务器、编写 Dockerfile 并正确配置 CUDA。管理这些事情非常麻烦，而且会破坏开发人员的整体体验。

为了避免这些问题，我决定使用一个名为 Truss 的简单开源工具。Truss 允许开发人员轻松打包支持 GPU 的模型，并在任何云环境中运行。它拥有大量强大的功能，让模型的容器化变得轻而易举：

• 开箱即支持 GPU。无需处理 CUDA。
• 自动创建 Dockerfile。无需自己编写。
• 生产就绪的 API 服务器
• 简单的 python 界面

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

创建桁架

创建或打开 python 虚拟环境（python 版本≥ 3.8），并安装以下依赖项：

pip install --upgrade truss

(可选）如果想从头开始创建 Truss 项目，可以运行以下命令：

truss init mistral-7b-tensort-llm

系统会提示你为模型命名。任何名称都可以，例如 Mistral 7B Tensort LLM。运行上述命令会自动生成部署 Truss 所需的文件。

这是 mistral-7b-tensorrt-llm-truss 的目录结构：

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

深入代码解释：

model.py包含执行的主要代码，因此让我们深入研究一下该文件。让我们先看看加载函数。

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
• 接下来，在加载函数中，我们使用 snapshot_download 函数下载编译后的模型。我的编译模型位于以下 repo id：htrivedi99/mistral-7b-v0.2-trtllm。如果你将编译后的模型上传到了其他地方，请相应更新此值。
• 然后，我们使用 model/utils.py 自带的 load_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}

• 预测函数接受一些模型输入，如 prompt、max_new_tokens、temperature 等。我们在函数顶部使用 request.pop 方法提取所有这些值。
• 接下来，我们使用 self.parse_input 辅助函数将提示格式化为 TensorRT LLM 所需的格式。
• 然后，我们调用 LLM 模型，使用 self.model.generate 函数生成输出。generate 函数接受多种参数，有助于控制 LLM 的输出。
• 我还添加了一些代码，通过生成生成器对象来启用流式处理。如果禁用了流式处理，标记生成器就会简单地解码 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

在 GKE 中部署模型

在本节中，我们将在谷歌 Kubernetes 引擎上部署模型。如果你还记得，在模型编译步骤中，我们在 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 内存）
• 安装谷歌管理的 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 开始运行，我们之前编写的加载函数就会被执行。你可以运行以下命令检查 pod 的日志：

kubectl logs <pod-name>

模型加载完成后，你会在 pod 日志中看到类似 Completed model.load() execution in 449234 ms 这样的信息。要通过 HTTP 向模型发送请求，我们需要对服务进行端口转发。你可以使用下面的命令来完成此操作：

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

我们终于可以开始向模型发送请求了！打开任何 Python 脚本并运行以下代码：

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."}

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 实现最先进的推理。我们涵盖了从编译 LLM 到在生产中部署模型的所有内容。

虽然 TensorRT LLM 比其他推理优化器更复杂，但其性能不言而喻。该工具提供了最先进的 LLM 优化，同时完全开源，专为商业用途而设计。该框架仍处于早期阶段，正在积极开发中。我们今天看到的性能只会在未来几年内不断提高。