TensorRT-LLM 提供易于使用的 Python API，用于定义大语言模型（LLM）并构建包含最先进优化的 TensorRT 引擎，以便在 NVIDIA GPU 上高效执行推理。

更多信息请参阅 TensorRT-LLM GitHub 仓库。

TensorRT-LLM 环境设置

由于 TensorRT-LLM 是用于在进程中与本地模型交互的 SDK，因此必须遵循一些环境步骤以确保可以使用 TensorRT-LLM。运行 TensorRT-LLM 需要 NVIDIA CUDA 12.2 或更高版本。

在本教程中，我们将展示如何将连接器与 GPT2 模型一起使用。为了获得最佳体验，我们建议按照官方 TensorRT-LLM Github 上的安装流程操作。

以下步骤展示如何为 x86_64 用户使用 TensorRT-LLM v0.8.0 设置模型。

1. 获取并启动基础 Docker 镜像环境。

docker run --rm --runtime=nvidia --gpus all --entrypoint /bin/bash -it nvidia/cuda:12.1.0-devel-ubuntu22.04

2. 安装依赖项，TensorRT-LLM 需要 Python 3.10

apt-get update && apt-get -y install python3.10 python3-pip openmpi-bin libopenmpi-dev git git-lfs wget

3. 安装 TensorRT-LLM 的最新稳定版本（对应发布分支）。我们使用的是 0.8.0 版本，但如需最新发布版本，请参阅官方发布页面。

pip3 install tensorrt_llm==0.8.0 -U --extra-index-url https://pypi.nvidia.com

4. 检查安装

python3 -c "import tensorrt_llm"

上述命令不应产生任何错误。

5. 本示例我们将使用 GPT2。GPT2 模型文件需要按照此处的说明通过脚本创建

   • 首先，在阶段 1 启动的容器内，克隆 TensorRT-LLM 仓库：

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

   • 使用以下命令安装 GPT2 模型的要求：

   cd TensorRT-LLM/examples/gpt/ && pip install -r requirements.txt

   • 下载 hf gpt2 模型

   rm -rf gpt2 && git clone https://huggingface.co/gpt2-medium gpt2
   cd gpt2
   rm pytorch_model.bin model.safetensors
   wget -q https://huggingface.co/gpt2-medium/resolve/main/pytorch_model.bin
   cd ..

   • 将权重从 HF Transformers 转换为 TensorRT-LLM 格式

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

   • 构建 TensorRT 引擎

   python3 build.py --model_dir=./c-model/gpt2/1-gpu --use_gpt_attention_plugin --remove_input_padding

6. 安装 llama-index-llms-nvidia-tensorrt 包

pip install llama-index-llms-nvidia-tensorrt

基本用法

使用提示词调用 complete

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=40,
)

resp = llm.complete("Who is Harry Potter?")
print(str(resp))

预期响应应类似于：

Harry Potter is a fictional character created by J.K. Rowling in her first novel, Harry Potter and the Philosopher's Stone. The character is a wizard who lives in the fictional town#