跳转到内容
LlamaIndex Python Documentation
TypeScript

英伟达 TensorRT-LLM

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

TensorRT-LLM GitHub

TensorRT-LLM 环境设置

Section titled “TensorRT-LLM Environment Setup”

由于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
  1. 安装依赖项，TensorRT-LLM 需要 Python 3.10
apt-get update && apt-get -y install python3.10 python3-pip openmpi-bin libopenmpi-dev git git-lfs wget
  1. 安装 TensorRT-LLM 的最新稳定版本（对应发布分支）。我们使用的是 0.8.0 版本，但如需获取最新版本，请参考官方发布页面。
pip3 install tensorrt_llm==0.8.0 -U --extra-index-url https://pypi.nvidia.com
  1. 检查安装
python3 -c "import tensorrt_llm"

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

  1. 对于此示例，我们将使用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://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

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

pip install llama-index-llms-nvidia-tensorrt

基本用法

Section titled “Basic Usage”

Call complete with a prompt

Section titled “Call complete with a prompt”
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#