使用Nvidia TensorRT-LLM进行本地大语言模型推理

在这篇文章中，我们将介绍如何使用Nvidia TensorRT-LLM进行本地大语言模型（LLM）推理。TensorRT-LLM为用户提供了一个易于使用的Python API，可以定义LLM并构建包含最先进优化的TensorRT引擎，以便在NVIDIA GPU上高效地执行推理。

环境设置

要使用TensorRT-LLM，您需要按照以下步骤进行环境设置：

1. 确保安装Nvidia Cuda 12.2或更高版本。

2. 使用以下命令通过pip安装TensorRT-LLM：

   pip3 install tensorrt_llm -U --extra-index-url https://pypi.nvidia.com
   

3. 准备Llama2模型文件。按照官方说明生成以下文件：

   • Llama_float16_tp1_rank0.engine: 包含嵌入模型权重的可执行操作图。
   • config.json: 包含模型的详细信息，如结构和精度，以及包含在引擎中的插件信息。
   • model.cache: 缓存模型编译中的一些计时和优化信息，使后续构建更快。

4. 创建模型目录并将上述文件移至该目录：

   mkdir model
   mv Llama_float16_tp1_rank0.engine config.json model.cache model/
   

5. 安装依赖包：

   %pip install llama-index-llms-nvidia-tensorrt
   !pip install tensorrt_llm==0.7.0 --extra-index-url https://pypi.nvidia.com --extra-index-url https://download.pytorch.org/whl/cu121
   

基本用法

以下是一个使用TensorRT-LLM进行推理的示例代码：

from llama_index.llms.nvidia_tensorrt import LocalTensorRTLLM

def completion_to_prompt(completion: str) -> str:
    """
    将补全结果转换为Llama2格式的提示。
    """
    return f"<s> [INST] {completion} [/INST] "

llm = LocalTensorRTLLM(
    model_path="./model",
    engine_name="llama_float16_tp1_rank0.engine",
    tokenizer_dir="meta-llama/Llama-2-13b-chat",
    completion_to_prompt=completion_to_prompt,
)

resp = llm.complete("Who is Paul Graham?")
print(str(resp))  # 使用中转API

可能遇到的错误

1. 环境依赖问题：确保已安装正确版本的CUDA和其他依赖包，否则可能会遇到兼容性问题。
2. 模型文件缺失：确保所有必要的模型文件都已正确生成并移动到指定目录。
3. 网络问题：如果下载依赖包或模型文件时遇到问题，可以尝试更换网络环境或使用国内镜像源。

如果你觉得这篇文章对你有帮助，请点赞，关注我的博客，谢谢!

参考资料：

• Nvidia TensorRT-LLM GitHub