ubuntu源码方式安装TensorRT-LLM推理框架(超详细)

记录安装过程和遇到的问题

写在前面：
建议先根据官方手册进行安装，遇到问题再看blog：
TensorRT官方文档

先安装docker

TensorRT-LLM 官方推荐使用 Docker 进行构建和运行
ubuntu安装docker官方手册
可以通过以下命令进行检测buildx是否可用（遇到的第一个安装错误，缺少buildx，重装docker解决了）

docker buildx version

获取 TensorRT-LLM 源码

apt-get update && apt-get -y install git git-lfs
git lfs install
git clone https://github.com/NVIDIA/TensorRT-LLM.git
cd TensorRT-LLM
git submodule update --init --recursive //更新并初始化子模块
git lfs pull

建立docker 镜像

可以搜索自己的GPU的计算能力，替换以下命令行的数字(能加速构建)
注：构建 TensorRT-LLM Docker 镜像需要约 63 GB 的磁盘空间

make -C docker release_build CUDA_ARCHS="86-real"

建立后可查看镜像：

创建并运行 Docker 容器

可以根据自己的环境修改Makefile文件

make -C docker release_run

如果希望以本地用户身份运行容器（而非 root），可以添加 LOCAL_USER=1 参数:

make -C docker release_run LOCAL_USER=1

遇到bug，因为NVIDIA Container Toolkit 未安装：
即使 NVIDIA 驱动已正确安装，Docker 仍需要通过 NVIDIA Container Toolkit 来支持 GPU 加速容器。
检查 NVIDIA Container Toolkit 是否安装, 运行以下命令检查 NVIDIA Container Toolkit 是否已安装：

docker run --rm --gpus all nvidia/cuda:12.6-base nvidia-smi

如果命令成功运行并显示 GPU 信息，则说明 NVIDIA Container Toolkit 已正确安装。
如果命令失败，则需要安装或重新配置 NVIDIA Container Toolkit。
安装 NVIDIA Container Toolkit, 按照以下步骤安装 NVIDIA Container Toolkit：

• 添加 NVIDIA 包仓库：

distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

• 更新包索引并安装 nvidia-docker2：

sudo apt-get update
sudo apt-get install -y nvidia-docker2

• 重启 Docker 服务：

sudo systemctl restart docker

在容器中构建 TensorRT-LLM

• 在容器中运行以下命令构建 TensorRT-LLM：

cd /code/tensorrt_llm
python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt

构建过程出现服务器崩溃，可以通过htop发现过载，导致内存溢出，修改Makefile文件（加 --cpus=“64” --memory=“64g”）重新创建容器，重新编译成功：

• 清理构建目录（可选）
  如果需要清理构建目录，可以添加 --clean 参数：

python3 ./scripts/build_wheel.py --clean --trt_root /usr/local/tensorrt

• 指定 GPU 架构（可选）
  如果需要限制支持的 GPU 架构，可以指定 --cuda_architectures 参数。例如，仅支持 Ampere 架构：

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

• 构建 C++ 基准测试（可选）
  如果需要构建 C++ 基准测试脚本，可以添加 --benchmarks 参数：

python3 ./scripts/build_wheel.py --benchmarks --trt_root /usr/local/tensorrt

安装 TensorRT-LLM

构建完成后，安装生成的 Wheel 包：

pip install ./build/tensorrt_llm*.whl

验证安装, 运行以下命令验证安装是否成功：

python3 -c "import tensorrt_llm; print(tensorrt_llm.__version__)"

跑一个离线推理模型

代码：

from tensorrt_llm import LLM, SamplingParams
def main():
    prompts = [
        "Hello, my name is",
        "The president of the United States is",
        "The capital of France is",
        "The future of AI is",
    ]
    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

    llm = LLM(model="/root/Meta-Llama-3-8B-Instruct")

    outputs = llm.generate(prompts, sampling_params)

    # Print the outputs.
    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text
        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
if __name__ == '__main__':
    main()

运行效果：

更多细节，访问TensorRT官方文档