TensorRT-LLM是一个由Nvidia设计的开源框架，用于在生产环境中进步大型语言模型的性能。该框架是基于 TensorRT 深度学习编译框架来构建、编译并执行计算图，并借鉴了许多 FastTransformer 中高效的 Kernels 实现，并且能够利用 NCCL 实现设施之间的通信。
尽管像vLLM和TGI这样的框架是加强推理的一个很好的终点，但它们不足一些优化，因而很难在生产中扩大它们。所以Nvidia在TensorRT的根底上有开发了TensorRT-LLM，像Anthropic, OpenAI, Anyscale等大公司曾经在应用这个框架为数百万用户提供LLM服务。
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官网的列表
TensorRT-LLM的益处
TensorRT LLM python包容许开发人员在不理解c++或CUDA的状况下以最高性能运行LLM。
分页注意力
大型语言模型须要大量内存来存储每个令牌的键和值。随着输出序列变长，这种内存应用会变得十分大。
通常状况下，序列的键和值必须间断存储。所以即便你在序列的内存调配中开释了空间，你也不能把这个空间用于其余序列。这会导致碎片化和节约。
分页注意力将键/值分成而不是间断的页，这样能够放在内存中的任何中央，如果您在两头开释一些分页，那么这些空间能够用于其余序列。
这能够避免碎片，并容许更高的内存利用率。在生成输入序列时，能够依据须要动静地调配和开释页面。
高效KV缓存
llm有数十亿个参数，这使得它们运行推理时速度迟缓且占用大量内存。KV缓存通过缓存LLM的层输入和激活来帮忙解决这个问题，因而它们不须要为每个推理从新计算。
上面是它的工作原理:
在推理期间，当LLM执行每一层时，输入将被缓存到具备惟一键的键值存储中。当后续推断应用雷同的层输出时，不是从新计算层，而是应用键检索缓存的输入。这防止了冗余计算，缩小了激活内存，进步了推理速度和内存效率。
上面咱们开始应用TensorRT-LLM部署一个模型
TensorRT-LLM部署教程
应用TensorRT-LLM部署模型首先就是要对模型进行编译，这里咱们将应用Mistral 7B instruction v0.2。编译阶段须要GPU，所以为了方便使用咱们间接在Colab上操作。
TensorRT LLM次要反对高端Nvidia gpu。所以咱们在Colab上抉择了A100 40GB GPU。
下载TensorRT-LLM git库。这个repo蕴含了编译模型所需的所有模块和脚本。
!git clone https://github.com/NVIDIA/TensorRT-LLM.git
%cd TensorRT-LLM/examples/llama而后装置所需的包
!pip install tensorrt_llm -U --pre --extra-index-url https://pypi.nvidia.com
!pip install huggingface_hub pynvml mpi4py
!pip install -r requirements.txt下载模型
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下一步就是应用trtllm-build命令编译模型。如果须要量化和其余优化，能够在这里指定参数。为了简略起见，我没有应用任何额定的优化。
!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 32256Mistral 7B instruction v0.2反对32K上下文长度。所以这里在max_input_length标记设置了上下文长度。
编译模型须要15-30分钟
模型编译实现后就能够间接应用了。咱们这里还要介绍一下模型部署的办法，有很多办法能够部署这个编译后的模型比方像FastAPI这样的简略工具，或者像triton推理服务器这样更简单的工具。
当应用像FastAPI这样的工具时，开发人员必须设置API服务器，编写Dockerfile，并正确配置CUDA，这外面蕴含了很多服务器后端的工作，有时候咱们并不相熟，所以这里咱们介绍一个简略的开源工具Truss。
Truss容许开发人员应用GPU打包他们的模型，并在任何云环境中运行它们。它有很多很棒的性能，使集成模型变得轻而易举。应用Truss的次要益处是，能够轻松地将具备GPU反对的模型容器化，并将其部署到任何云环境中。
装置:
pip install --upgrade truss如果从头开始创立Truss我的项目，你能够运行上面的命令:
truss init mistral-7b-tensort-llmmistral-7b-tensort-llm是咱们我的项目的名称，能够轻易编写。运行下面的命令会主动生成部署Truss所需的文件。
上面是mistral-7b- tensort -llm-truss的目录构造:
├── mistral-7b-tensorrt-llm-truss
│ ├── config.yaml
│ ├── model
│ │ ├── __init__.py
│ │ └── model.py
| | └── utils.py
| ├── requirements.txt以下是上述文件的疾速介绍:
1、config.yaml用于为模型设置各种配置，包含其资源、依赖项、环境变量等。在这里，咱们能够指定模型名称、要装置的Python依赖项以及要装置的零碎包。
2、model/model.py是Truss的外围。它蕴含将在Truss服务器上执行的Python代码。在model.py中有两个次要办法:load()和predict()。
load办法是咱们从hugs face下载编译模型并初始化TensorRT LLM的中央；predict办法接管HTTP申请并调用模型。
3、model/utils.py蕴含model.py文件的一些辅助函数。utils.py文件不是咱们本人编写的，能够间接从TensorRT LLM存储库中获取的。
4、含运行编译模型所需的Python依赖项，truss会应用它来初始化咱们的环境。
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函数下载编译后的模型；而后应用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}predict函数承受一些模型输出，如提醒、max_new_tokens、温度等。咱们应用申请在函数的顶部提取所有这些值。调用LLM模型来应用self.model.generate函数生成输入。generate函数承受各种参数，帮忙管制LLM的输入。
为了在云中运行咱们的模型，还须要将其容器化。Truss会负责为咱们创立Dockerfile并打包所有内容，所以咱们不须要做太多事件。
在mistral-7b- tensort -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- tensort -llm-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这些docker的配置文件就是truss为咱们主动生成好的，咱们上面简略的介绍一下看k8s的部署，我不会深刻探讨如何设置GKE集群，因为这不在本文的探讨范畴之内。
创立以下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日志中看到相似Completed model.load()的执行工夫为449234毫秒。要通过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的减速推理还是很显著的
总结
在这篇文章中，咱们演示了如何应用TensorRT LLM实现模型减速推理，文章内容涵盖了从编译LLM到在生产中部署模型的所有内容。
尽管TensorRT LLM比其余推理优化器更简单，但性能进步也是非常明显。尽管该框架仍处于晚期阶段，然而能够提供目前最先进的LLM优化。并且它是齐全开源的能够商业化，我置信TensorRT LLM当前还会有更大的倒退，因为毕竟是NVIDIA本人的产品.
TensorRT-LLM代码：
https://avoid.overfit.cn/post/22b19ff044984de69da655a67721cff3
作者：Het Trivedi
发表回复