TensorRT-LLM 在 PyPi 上的安装与配置指南

随着大规模语言模型的广泛应用，TensorRT-LLM 推理框架在降低模型推理延迟和提升吞吐量方面表现出色。本文将详述如何通过 PyPi 安装和配置 TensorRT-LLM，帮助用户在生产环境中高效部署大模型。

TensorRT-LLM 简介

TensorRT-LLM 是一个用于大规模模型推理的框架，能够显著减少推理时间并提高吞吐量。其支持多种量化技术，包括 FP16、INT8 和 SmoothQuant。通过这些技术，TensorRT-LLM 可以在不显著影响模型精度的前提下，降低计算资源消耗。

TensorRT-LLM 的优势

TensorRT-LLM 可以自动优化和转换深度学习模型，使得在 GPU 上运行更加高效。这不仅减少了开发人员的工作量，还能最大化硬件性能利用率。此外，TensorRT-LLM 的插件机制允许用户自定义算子，进一步提升推理性能。

环境准备与安装

在开始安装 TensorRT-LLM 之前，需要确保系统满足以下基本要求：

CUDA 版本：12.2 或更高
PyTorch 镜像：推荐使用 nvcr.io/nvidia/pytorch:23.10-py3

由于某些服务器无法访问外网，用户需提前准备好相关镜像和安装包。

Docker 环境搭建

Docker 是安装 TensorRT-LLM 的推荐方式。首先，拉取并运行 PyTorch 镜像，进入 Docker 容器。

    docker run -dt --name tensorrt_llm_container 
    --restart=always 
    --gpus all 
    --network=host 
    --shm-size=4g 
    -m 64G 
    -v /home/user/workspace:/workspace 
    -w /workspace 
    nvcr.io/nvidia/pytorch:23.10-py3 
    /bin/bash

    docker exec -it tensorrt_llm_container bash

安装 PyTorch 和 TensorRT

完成 Docker 环境搭建后，接下来安装必要的 Python 包，包括 PyTorch 和 TensorRT。

    pip uninstall -y tensorrt
    pip uninstall -y torch-tensorrt

    pip install mpi4py -i http://nexus3.yourcompany.com/repository/pypi/simple --trusted-host nexus3.yourcompany.com
    pip install polygraphy-0.48.1-py2.py3-none-any.whl -i http://nexus3.yourcompany.com/repository/pypi/simple --trusted-host nexus3.yourcompany.com

    pip install torch==2.1.0 -i http://nexus3.yourcompany.com/repository/pypi/simple --trusted-host nexus3.yourcompany.com
    pip uninstall transformer-engine

接着，安装 TensorRT。

    tar -xf /tmp/TensorRT.tar -C /usr/local/
    mv /usr/local/TensorRT-9.1.0.4 /usr/local/tensorrt
    pip install /usr/local/tensorrt/python/tensorrt-*-cp310-*.whl -i http://nexus3.yourcompany.com/repository/pypi/simple --trusted-host nexus3.yourcompany.com

配置环境变量

在 /etc/profile 中配置环境变量，以便系统正确加载 TensorRT。

    export LD_LIBRARY_PATH=/usr/local/tensorrt/lib:${LD_LIBRARY_PATH}

构建 TensorRT-LLM

为了顺利构建 TensorRT-LLM，需要在 Docker 容器中执行以下命令。

    python3 ./scripts/build_wheel.py --clean --trt_root /usr/local/tensorrt --cuda_architectures "80-real"

在构建过程中，可能需要对配置文件进行修改，例如更改 pip 源地址或 git 远程仓库地址。

基于 Bloom 模型的实践

在完成环境配置后，本文将以 Bloom 模型为例，讲解如何在 TensorRT-LLM 中进行模型量化和推理。

Bloom 模型简介

Bloom 是一个大规模语言模型，支持多种量化方式。其主要功能文件包括：

build.py：用于构建 TensorRT 引擎。
run.py：用于模型推理。
summarize.py：用于总结数据集的文章。
hf_bloom_convert.py：用于模型格式转换。

下载与准备

首先，下载 Bloom 模型及相关数据集。

    rm -rf /workspace/model/bloomz-3b
    mkdir -p /workspace/model/bloomz-3b && git clone https://huggingface.co/bigscience/bloomz-3b /workspace/model/bloomz-3b

数据集下载链接：

构建 TensorRT 引擎

使用 HF 上 Bloom 的 checkpoint 构建 TensorRT 引擎。未指定 checkpoint 目录时，将使用虚拟权重构建引擎。

FP16 构建引擎

以下命令用于构建基于单 GPU 的 FP16 精度引擎。

    python build.py --model_dir /workspace/model/bloomz-3b 
                    --dtype float16 
                    --use_gemm_plugin float16 
                    --use_gpt_attention_plugin float16 
                    --output_dir /workspace/model/bloomz-3b_trt_engines/fp16/1-gpu/

模型推理

配置完引擎后，接下来运行模型进行推理，并使用 Rouge 指标评估模型。

    python summarize.py --test_trt_llm 
                        --hf_model_location /workspace/model/bloomz-3b 
                        --data_type fp16 
                        --engine_dir /workspace/model/bloomz-3b_trt_engines/fp16/1-gpu/

常见问题解答 (FAQ)

FAQ

问：TensorRT-LLM 支持哪些量化方式？
- 答：TensorRT-LLM 支持 FP16、INT8、SmoothQuant 等多种量化方式，适用于不同的推理需求。
问：如何解决 Docker 环境无法访问外网的问题？
- 答：可以提前下载所需的镜像和安装包，并通过本地挂载的方式在 Docker 中使用。
问：安装 TensorRT 时需要注意哪些事项？
- 答：确保 CUDA 版本和 TensorRT 版本的兼容性，并正确配置环境变量 LD_LIBRARY_PATH。
问：如何优化 TensorRT-LLM 的推理性能？
- 答：可以通过调整量化参数和并行配置来优化推理性能，具体方法可参考官方文档。
问：在哪里可以获取更多关于 TensorRT-LLM 的信息？
- 答：可以访问 GitHub 仓库 llm-action 获取更多的博客和代码示例。