Nvidia TensorRT-LLM¶

TensorRT-LLM 为用户提供易于使用的 Python API，用于定义大语言模型 (LLMs) 并通过包含前沿优化的 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
   

   • 下载 huggingface 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#