Nvidia tensorrt

Nvidia TensorRT-LLM是一种用于低延迟推理的深度学习模型优化工具。TensorRT-LLM使用Nvidia的TensorRT推理引擎，通过减少模型推理的延迟来提高性能。它还提供了一些优化技术，如层次化的量化和剪枝，以进一步优化模型。TensorRT-LLM旨在帮助开发人员在边缘设备上部署实时的深度学习推理应用程序。

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

TensorRT-LLM Github

TensorRT-LLM环境设置¶

由于TensorRT-LLM是一个用于与本地模型进行进程交互的SDK，因此必须遵循一些环境设置步骤，以确保可以使用TensorRT-LLM设置。请注意，目前需要Nvidia Cuda 12.2或更高版本才能运行TensorRT-LLM。

在本教程中，我们将展示如何将连接器与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("谁是哈利波特？")
print(str(resp))

哈利·波特是J.K.罗琳在她的第一部小说《哈利·波特与魔法石》中创造的虚构人物。这个角色是一个生活在虚构小镇的巫师。