Nvidia TensorRT-LLM¶

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 模型文件。

   • 首先，在我们第一阶段启动的容器内，克隆 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://hugging-face.cn/gpt2-medium gpt2
   cd gpt2
   rm pytorch_model.bin model.safetensors
   wget -q https://hugging-face.cn/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#