Nvidia TensorRT-LLM是一种用于低延迟推理的深度学习模型优化工具。TensorRT-LLM使用Nvidia的TensorRT推理引擎,通过减少模型推理的延迟来提高性能。它还提供了一些优化技术,如层次化的量化和剪枝,以进一步优化模型。TensorRT-LLM旨在帮助开发人员在边缘设备上部署实时的深度学习推理应用程序。
TensorRT-LLM为用户提供了一个易于使用的Python API,用于定义大型语言模型(LLMs)并构建包含最先进优化的TensorRT引擎,以在NVIDIA GPU上高效进行推理。
由于TensorRT-LLM是一个用于与本地模型进行进程交互的SDK,因此必须遵循一些环境设置步骤,以确保可以使用TensorRT-LLM设置。请注意,目前需要Nvidia Cuda 12.2或更高版本才能运行TensorRT-LLM。
在本教程中,我们将展示如何将连接器与GPT2模型一起使用。 为了获得最佳体验,我们建议按照官方TensorRT-LLM Github上的安装过程进行操作。
以下步骤展示了如何为x86_64用户设置您的模型与TensorRT-LLM v0.8.0。
docker run --rm --runtime=nvidia --gpus all --entrypoint /bin/bash -it nvidia/cuda:12.1.0-devel-ubuntu22.04
apt-get update && apt-get -y install python3.10 python3-pip openmpi-bin libopenmpi-dev git git-lfs wget
pip3 install tensorrt_llm==0.8.0 -U --extra-index-url https://pypi.nvidia.com
python3 -c "import tensorrt_llm"
上述命令不应产生任何错误。
对于本示例,我们将使用GPT2。GPT2模型文件需要通过脚本按照这里的说明创建。
git clone --branch v0.8.0 https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM/examples/gpt/ && pip install -r requirements.txt
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 ..
python3 hf_gpt_convert.py -i gpt2 -o ./c-model/gpt2 --tensor-parallelism 1 --storage-type float16
python3 build.py --model_dir=./c-model/gpt2/1-gpu --use_gpt_attention_plugin --remove_input_padding
安装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.罗琳在她的第一部小说《哈利·波特与魔法石》中创造的虚构人物。这个角色是一个生活在虚构小镇的巫师。