在之前的讨论中，我们探讨了 NVIDIA TensorRT 在加速深度学习推理方面的能力。今天，我们将这一探索扩展到 TensorRT-LLM，这是一个专门的开源库，旨在优化和加速 NVIDIA GPU 上的大型语言模型 (LLM) 推理。

TensorRT-LLM 以 TensorRT 的性能优势为基础，增强了 LLM 部署，使 Meta 的 Llama 和 Mistral 等强大的模型不仅可以在数据中心高效运行，还可以在配备 RTX GPU 的 Windows PC 上高效运行。该库通过飞行批处理和张量并行等先进技术，将推理性能提高 8 倍，从而可以在本地运行尖端的 AI 应用程序。

通过利用 TensorRT-LLM，开发人员可以减少对云基础设施的依赖，从而节省成本并提高数据隐私，同时实现一流的性能。 TensorRT-LLM 即将推出更新，旨在提供与流行模型和工具的更大集成，从而让高性能 LLM 功能可在各种平台上使用。

1、安装 TensorRT-LLM

在这里，我们首先更新系统包并安装必要的依赖项，如 Python 3.10、pip、OpenMPI、Git 和 wget。然后，我们通过 pip 安装 TensorRT-LLM（版本 0.8.0），从 NVIDIA 的 PyPI 存储库获取它，以确保与 Tensor Core GPU 兼容。

apt-get update && apt-get -y install python3.10 python3-pip openmpi-bin libopenmpi-dev git git-lfs wget
pip3 install tensorrt_llm==0.8.0 -U --extra-index-url https://pypi.nvidia.com

安装 TensorRT-LLM 后，我们通过导入库快速验证设置。如果安装成功，输出应确认 TensorRT-LLM 的版本为 0.8.0。

python3 -c "import tensorrt_llm"

我们应该得到： [TensorRT-LLM] TensorRT-LLM Version: 0.8.0

2、安装要求和克隆模型示例

接下来，我们克隆 0.8.0 版本的 TensorRT-LLM 存储库并导航到受支持模型（例如 GPT-2）的示例目录。

虽然这里使用 GPT-2 作为示例，但 TensorRT-LLM 支持各种模型，你可以用你选择的任何其他受支持模型替换它。然后，我们从 requirements.txt 文件安装必要的依赖项，为在 TensorRT-LLM 中使用所选模型准备环境。

git clone --branch v0.8.0 https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM/examples/gpt

pip install -r requirements.txt

在这里，我们将 Hugging Face 中的 GPT-2 模型存储库直接克隆到 TensorRT-LLM 示例目录中。然后，我们清理不必要的文件，只保留 pytorch_model.bin 文件（如果有），该文件是转换为 TensorRT-LLM 格式所必需的。

cd TensorRT-LLM/examples/gpt/
rm -rf gpt2 && git clone https://huggingface.co/gpt2 gpt2
cd gpt2
rm model.safetensors
cd ..

3、处理缺失的 .bin 模型

如果 GPT-2 模型存储库缺少 pytorch_model.bin，我们会使用脚本来生成它。该脚本使用 Transformers 库加载模型和 tokenizer，并以 .bin 格式保存模型状态。它还会删除 .safetensors 文件（如果存在），以确保与转换过程兼容。

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os
import glob

def save_model_as_bin(model_name, save_directory):
    # Load the model and tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)

    # Create save directory if it doesn't exist
    os.makedirs(save_directory, exist_ok=True)

    # Save the model and tokenizer
    model.save_pretrained(save_directory)
    tokenizer.save_pretrained(save_directory)

    # Save the model in .bin format
    torch.save(model.state_dict(), os.path.join(save_directory, "pytorch_model.bin"))

    # Remove .safetensors files if they exist
    for safetensors_file in glob.glob(os.path.join(save_directory, "*.safetensors")):
        os.remove(safetensors_file)

# Usage
save_model_as_bin("model_id", "TensorRT-LLM/examples/....")

4、将权重转换为 TensorRT-LLM 格式

转换步骤将 Hugging Face 模型权重转换为 TensorRT-LLM 所需的格式。这涉及指定输入目录、输出路径、张量并行度（本例中为 1）以及使用 float16 进行存储，以优化推理性能。

python3 TensorRT-LLM/examples/gpt/hf_gpt_convert.py -i TensorRT-LLM/examples/gpt/gpt2 -o ./c-model/gpt2 --tensor-parallelism 1 --storage-type float16

5、构建 TensorRT 引擎

现在，我们构建 TensorRT 引擎，它是模型的高度优化表示，可用于快速推理。此脚本使用转换后的模型文件并指定选项，例如使用 GPT 注意力插件和删除输入填充，从而进一步提高性能。

python3 TensorRT-LLM/examples/gpt/build.py --model_dir=./c-model/gpt2/1-gpu --use_gpt_attention_plugin --remove_input_padding

6、使用 Llama-index 进行推理

我们可以使用 Llama-index 库对优化的 TensorRT 引擎进行推理。

pip install llama-index llama-index-llms-nvidia-tensorrt

安装必要的 Llama-index 包后，代码会实例化 LocalTensorRTLLM 对象，提供引擎、标记器和其他参数的路径。

from llama_index.llms.nvidia_tensorrt import LocalTensorRTLLM

llm = LocalTensorRTLLM(
    model_path="./engine_outputs",
    engine_name="gpt_float16_tp1_rank0.engine",
    tokenizer_dir="gpt2",
    max_new_tokens=15,
)

然后演示了一个简单的完成任务，用“once upon a time”提示模型并打印生成的文本……

resp = llm.complete("once upon a time")
print(str(resp))

7、使用 Triton Inference Server 部署

在这里，我们首先克隆特定于版本 0.9.0 的 TensorRT-LLM 后端存储库，并将模型文件从 c-model/gpt2/1-gpu 目录复制到 Triton 后端模型目录中。

git clone -b v0.9.0 https://github.com/triton-inference-server/tensorrtllm_backend.git
cp c-model/gpt2/1-gpu/* tensorrtllm_backend/all_models/inflight_batcher_llm/tensorrt_llm/1

这将为后端准备部署所需的模型文件。

8、更新一些配置

接下来，我们通过为不同的组件创建配置文件来配置 TensorRT-LLM 以支持飞行中批处理。这包括使用 fill_template.py 脚本填写预处理、后处理和主模型配置的模板。飞行中批处理可以一起处理请求，以提高吞吐量并减少延迟。

HF_LLAMA_MODEL=/content/TensorRT-LLM/examples/gpt/gpt2
ENGINE_PATH=tensorrtllm_backend/all_models/inflight_batcher_llm/tensorrt_llm/1
python3 tensorrtllm_backend/tools/fill_template.py -i tensorrtllm_backend/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 tensorrtllm_backend/tools/fill_template.py -i tensorrtllm_backend/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 tensorrtllm_backend/tools/fill_template.py -i tensorrtllm_backend/all_models/inflight_batcher_llm/tensorrt_llm_bls/config.pbtxt \
        triton_max_batch_size:64,decoupled_mode:True,bls_instance_count:1,accumulate_tokens:False

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

python3 tensorrtllm_backend/tools/fill_template.py -i tensorrtllm_backend/all_models/inflight_batcher_llm/tensorrt_llm/config.pbtxt \
        triton_max_batch_size:64,decoupled_mode:True,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

指定 triton_max_batch_size、 decoupled_mode（用于流式传输）和 batching_strategy 等配置设置以优化性能。

9、提供模型服务

为了提供模型服务，我们在 Docker 容器中运行 Triton 推理服务器。 docker run 命令设置具有 GPU 支持的服务器，将当前目录挂载到容器，并指定工作目录。

docker run -it --rm --gpus all --network host --shm-size=1g -v $(pwd):/tensorrtllm_backend --workdir /tensorrtllm_backend nvcr.io/nvidia/tritonserver:24.03-trtllm-python-py3

我们可以使用提供的 Python 脚本启动 Triton 服务器，设置模型存储库路径和世界大小以配置服务器进行模型推理。

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

10、测试部署

在这里，我们通过向 Triton 服务器的 API 端点发送请求来测试已部署的模型。Python 脚本使用输入文本和参数构建有效负载，向模型端点发送 POST 请求，并打印响应。

import requests

url = "http://localhost:8000/v2/models/ensemble/generate"
payload = {
    "text_input": "Say this is a test.",
    "parameters": {
        "max_tokens": 128,
        "stop_words": [""]
    }
}

response = requests.post(url, json=payload)

if response.status_code == 200:
    print("Response:", response.json())
else:
    print(f"Error {response.status_code}: {response.text}")

这验证了模型是否正确提供并可以按预期生成响应。

流式传输测试：

import requests

url = "http://localhost:8000/v2/models/ensemble/generate_stream"
payload = {
    "text_input": "Tell me a short joke about llamas",
    "parameters": {
        "max_tokens": 128,
        "stop_words": [""],
        "stream": True
    }
}

# Send a POST request with streaming enabled
response = requests.post(url, json=payload, stream=True)

if response.status_code == 200:
    # Process the streaming response
    for line in response.iter_lines():
        if line:
            # Print or process each line of the stream
            print(line.decode('utf-8'))
else:
    print(f"Error {response.status_code}: {response.text}")

11、设置 OpenAI 兼容 API

为了将 TensorRT-LLM 与 OpenAI 兼容 API 集成，我们首先克隆 openai_trtllm 存储库，它在 OpenAI 的 API 格式和 Triton 推理服务器之间架起了一座桥梁。

git clone --recursive https://github.com/npuichigo/openai_trtllm.git

确保已安装 Rust，然后使用 Cargo 构建项目。

cargo run --release

这会生成运行 API 服务器所需的二进制文件。

12、运行与 OpenAI 兼容的 API

构建源代码后，我们运行 openai_trtllm 二进制文件来启动 API 服务器。我们使用主机、端口、Triton 端点等选项配置服务器……

./target/release/openai_trtllm --help
Usage: openai_trtllm [OPTIONS]

Options:
  -H, --host <HOST>
          Host to bind to [default: 0.0.0.0]
  -p, --port <PORT>
          Port to bind to [default: 3000]
  -t, --triton-endpoint <TRITON_ENDPOINT>
          Triton gRPC endpoint [default: http://localhost:8001]
  -o, --otlp-endpoint <OTLP_ENDPOINT>
          Endpoint of OpenTelemetry collector
      --history-template <HISTORY_TEMPLATE>
          Template for converting OpenAI message history to prompt
      --history-template-file <HISTORY_TEMPLATE_FILE>
          File containing the history template string
      --api-key <API_KEY>
          Api Key to access the server
  -h, --help
          Print help

openai_trtllm — help 命令显示可用于自定义服务器行为的选项。

测试与 OpenAI 兼容的 API：

./target/release/openai_trtllm

Openai 客户端：

import pprint

from openai import OpenAI

client = OpenAI(base_url="http://localhost:3000/v1", api_key="test")

result = client.completions.create(
    model="ensemble",
    prompt="Say this is a test",
)
pprint.pprint(result)

流式测试：

from sys import stdout

from openai import OpenAI

client = OpenAI(base_url="http://localhost:3000/v1", api_key="test")

response = client.completions.create(
    model="ensemble",
    prompt="This is a story of a hero who went",
    stream=True,
    max_tokens=50,
)
for event in response:
    if not isinstance(event, dict):
        event = event.model_dump()
    event_text = event["choices"][0]["text"]
    stdout.write(event_text)
    stdout.flush()

Langchain ChatOpenAI :

from langchain.chat_models import ChatOpenAI
from langchain.schema.messages import HumanMessage, SystemMessage


chat = ChatOpenAI(openai_api_base="http://localhost:3000/v1",
                  openai_api_key="test", model_name="ensemble",
                  max_tokens=100)

messages = [
    SystemMessage(content="You're a helpful assistant"),
    HumanMessage(content="What is the purpose of model regularization?"),
]

result = chat.invoke(messages)
print(result.content)

13、结束语

使用 Triton Inference Server 部署 TensorRT-LLM 并设置与 OpenAI 兼容的 API 可实现高效、高性能的推理大型语言模型。通过利用实时批处理和优化模型配置，你可以显著改善吞吐量和延迟。与 Triton Inference Server 的集成允许无缝模型服务，而与 OpenAI 兼容的 API 有助于轻松与现有应用程序集成。此设置不仅提高了性能，而且还提供了在本地或生产环境中部署和与高级 AI 模型交互的灵活性。

优化愉快！！

原文链接：大模型推理加速与服务优化 - 汇智网