2. 环境搭建

核心原理:为什么环境搭建是推理加速的第一道门槛

TensorRT-LLM 并非一个独立的推理框架,而是建立在 NVIDIA GPU 驱动、CUDA 工具包、cuDNN、TensorRT 库 以及 Python 生态 之上的一个“推理加速层”。它的核心工作方式是将大语言模型的计算图在离线阶段编译为针对特定 GPU 架构优化的引擎(Engine),并在运行时通过 CUDA Graph、KV Cache 管理、PagedAttention 等技术实现低延迟推理。因此,任何一层基础组件的版本不匹配,都会导致编译失败或运行时崩溃。

理解环境搭建的关键在于把握 “版本兼容矩阵” 这一概念。TensorRT-LLM 的发布版本与 CUDA 版本、cuDNN 版本、TensorRT 版本以及 GPU 架构(如 Ampere、Ada、Hopper)之间存在严格的对应关系。例如,TensorRT-LLM 0.9.0 版本要求 CUDA 12.2 及以上,而某些早期版本仅支持 CUDA 11.8。如果盲目安装最新版 CUDA,可能导致 TensorRT-LLM 的预编译 Wheel 包无法识别,从而被迫从源码编译,显著增加安装复杂度。

此外,GPU 驱动与 CUDA 的关系也常被误解:GPU 驱动是内核级组件,CUDA 是用户级工具包。驱动决定了 GPU 能支持的最高 CUDA 版本,而 CUDA 工具包则提供编译器和运行时库。安装 TensorRT-LLM 时,我们通常只需要确保驱动版本足够新(支持所需 CUDA 版本),而无需在系统层面安装完整 CUDA,因为 TensorRT-LLM 的 Python Wheel 包会自带所需的 CUDA 运行时库(通过 pip 依赖自动拉取)。

版本兼容性速查表(以 TensorRT-LLM 0.10.0 为例)

组件推荐版本说明
GPU 驱动≥ 535.104.05支持 CUDA 12.3 的最低驱动版本
CUDA 工具包12.2 / 12.3用于编译自定义算子,运行时由 Wheel 包提供
cuDNN8.9.5+深度学习原语库,TensorRT-LLM 依赖其卷积与注意力实现
TensorRT9.2.0引擎编译核心库,Wheel 包会自动安装
Python3.8 - 3.11推荐 3.10,兼容性最佳
GPU 架构Ampere / Ada / HopperV100(Volta)不支持,需旧版本
TensorRT-LLM 环境依赖栈与安装验证流程 GPU 驱动(内核级) nvidia-smi 验证 | 版本 ≥ 535 | 决定 GPU 支持的最高 CUDA 版本 CUDA 工具包(用户级) nvcc --version 验证 | 版本 12.2/12.3 | 提供 nvcc 编译器与运行时库 TensorRT 库 引擎编译核心 | 版本 9.2+ cuDNN 库 深度学习原语 | 版本 8.9.5+ TensorRT-LLM Python 包 pip install tensorrt-llm | 自动拉取 TensorRT/cuDNN/CUDA 运行时依赖 安装验证 python -c "import tensorrt_llm; print(tensorrt_llm.__version__)" | 运行示例模型 ⚠ 常见错误 驱动版本过旧 CUDA 与驱动不匹配 Python 版本过高

方法流程:从零到可用的四步安装法

环境搭建的核心方法论可以概括为 “先驱动、再 CUDA、后 Python 包、终验证” 四个阶段。每一步都有明确的验证命令,确保问题在最早阶段暴露,避免后续排查困难。

第一步:GPU 驱动安装与验证

驱动是硬件与软件之间的桥梁。在 Ubuntu 22.04 系统上,推荐使用 nvidia-driver-535 或更高版本。安装前,先检查系统是否已有 NVIDIA 显卡:

# 检查 GPU 硬件
lspci | grep -i nvidia

# 安装驱动(Ubuntu 22.04 示例)
sudo apt update
sudo apt install -y nvidia-driver-535

# 重启后验证驱动
nvidia-smi

如果 nvidia-smi 能正常输出 GPU 信息(包括驱动版本、显存使用率),说明驱动安装成功。注意:驱动安装后必须重启系统,否则模块无法加载。

第二步:CUDA 工具包安装(可选但推荐)

虽然 TensorRT-LLM 的 Wheel 包自带 CUDA 运行时,但为了编译自定义算子或调试,建议安装 CUDA 工具包。推荐使用 runfile 安装方式,避免污染系统目录:

# 下载 CUDA 12.2 runfile(约 3GB)
wget https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda_12.2.0_535.54.03_linux.run

# 执行安装(仅安装工具包,不覆盖驱动)
sudo sh cuda_12.2.0_535.54.03_linux.run --toolkit --silent --override

# 配置环境变量
echo 'export PATH=/usr/local/cuda-12.2/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.2/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

# 验证 CUDA 编译器
nvcc --version

如果输出 CUDA 12.2 版本信息,说明 CUDA 工具包安装成功。此时可以运行 nvidia-smi 确认驱动版本仍与 CUDA 12.2 兼容(驱动版本 ≥ 535 即可)。

第三步:安装 TensorRT-LLM Python 包

TensorRT-LLM 提供了预编译的 Wheel 包,支持 pip 直接安装。推荐使用 虚拟环境 隔离依赖,避免与系统 Python 包冲突:

# 创建 Python 3.10 虚拟环境
python3.10 -m venv trtllm_env
source trtllm_env/bin/activate

# 安装 TensorRT-LLM(以 0.10.0 版本为例)
pip install tensorrt-llm==0.10.0

# 验证安装
python -c "import tensorrt_llm; print('TensorRT-LLM 版本:', tensorrt_llm.__version__)"

如果输出版本号且无报错,说明核心安装成功。此时 pip 会自动安装 TensorRT、cuDNN、CUDA 运行时等依赖,无需手动干预。

第四步:运行示例模型验证完整链路

仅导入包成功还不够,必须实际运行一个模型推理,验证 GPU 计算链路是否通畅。TensorRT-LLM 提供了示例脚本:

# 下载示例模型(以 GPT-2 为例)
cd examples/gpt
python convert_checkpoint.py --model-dir gpt2 \
    --output-dir ./models/gpt2-ckpt \
    --model-type gpt

# 构建 TensorRT 引擎
trtllm-build --checkpoint_dir ./models/gpt2-ckpt \
    --output_dir ./models/gpt2-engine \
    --gemm_plugin float16

# 运行推理
python run.py --engine_dir ./models/gpt2-engine \
    --tokenizer_dir gpt2 \
    --max_output_len 50

如果模型能正常生成文本,说明从驱动到 CUDA 再到 TensorRT-LLM 的完整链路已打通。此时环境搭建正式完成。

案例/场景:Docker 容器化部署的推荐实践

在实际生产环境中,我们强烈推荐使用 NVIDIA NGC 容器 来规避环境依赖问题。NGC 容器已预装匹配的驱动、CUDA、cuDNN 和 TensorRT-LLM,只需安装 NVIDIA Container Toolkit 即可直接使用:

# 安装 NVIDIA Container Toolkit
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
sudo apt update && sudo apt install -y nvidia-container-toolkit
sudo systemctl restart docker

# 拉取 TensorRT-LLM NGC 容器
docker pull nvcr.io/nvidia/tensorrt-llm:24.05

# 启动容器并挂载 GPU
docker run --gpus all -it --rm \
    -v /path/to/models:/models \
    nvcr.io/nvidia/tensorrt-llm:24.05

进入容器后,直接运行 python -c "import tensorrt_llm" 即可验证。容器化方案的优势在于:版本一致性(团队所有成员使用相同镜像)、快速回滚(镜像 Tag 即版本)、隔离性(不影响宿主机环境)。

实操建议:如果使用 Docker 容器,宿主机只需安装 GPU 驱动和 NVIDIA Container Toolkit,无需安装 CUDA 工具包。容器内已包含全部依赖,这能减少 80% 的环境配置问题。

误区/复盘:环境搭建中的常见陷阱

根据社区反馈和实际教学经验,以下四个误区最容易导致安装失败或运行异常,需要特别注意:

  • 误区一:驱动版本过旧,却强行安装新版 CUDA。 驱动与 CUDA 存在“向下兼容”关系:驱动版本决定了支持的最高 CUDA 版本。如果驱动为 470 系列,最高仅支持 CUDA 11.4,安装 CUDA 12.x 会导致运行时错误。解决方法是先升级驱动,再安装 CUDA。
  • 误区二:使用系统 Python 直接安装,导致依赖冲突。 系统 Python 通常被多个系统工具依赖,pip 安装 TensorRT-LLM 可能覆盖关键库(如 numpy、protobuf),导致系统崩溃。务必使用虚拟环境或 Conda 环境。
  • 误区三:忽略 cuDNN 版本要求。 TensorRT-LLM 对 cuDNN 版本有严格限制,版本过旧会导致注意力算子编译失败。虽然 Wheel 包会自动安装 cuDNN,但如果系统存在旧版本且 LD_LIBRARY_PATH 优先级更高,会引发“找不到 cuDNN 符号”的错误。
  • 误区四:在 WSL2 中直接安装驱动。 WSL2 不支持在 Linux 内部安装 NVIDIA 驱动,必须使用 Windows 宿主机的驱动,并通过 WSL 的 CUDA 转发机制。正确做法是在 Windows 侧安装驱动,WSL 内仅安装 CUDA 工具包。

排障提示:如果遇到 ImportError: libcudnn.so.8: cannot open shared object file,说明 cuDNN 库路径未正确配置。执行 find / -name "libcudnn*" 2>/dev/null 找到库文件位置,然后将其所在目录添加到 LD_LIBRARY_PATH。若问题依旧,可尝试在虚拟环境中重新安装 nvidia-cudnn-cu12 包。