在AutoDL上编译TensorRT-LLM

背景

想要在B站简单分享一下TensorRT-LLM适配经验，顺便复习相关理论。
考虑到可能有大量用户没有显卡，但是他们又想来体验TensorRT-LLM，所以准备用源码来编译一波TensorRT-LLM，然后制作成公开镜像在AutoDL上面进行分享。
该镜像只支持RTX30系(Ampere)/40系（ADA)，V100/A100/A10/H100等卡，不支持20系/T4/P40等显卡。
我的仓库：https://github.com/Tlntin/Qwen-TensorRT-LLM

准备工作

观察一下现有TensorRT-LLM容器中的系统配置。
通过观察该文件

bash

ARG BASE_IMAGE=nvcr.io/nvidia/pytorch
ARG BASE_TAG=23.10-py3

可以看出目前基础容器是nvcr.io/nvidia/pytorch:23.10-py3
再去官网看看这个容器里面的详细信息，官网
下面是23.10-py3对应的系统组件版本
组件版本
Container OS Ubuntu 22.04
CUDA NVIDIA CUDA 12.2.2
cuDNN 8.9.5
Python 3.10.12
TensorRT在编译的时候会另外安装，根据下面这个文件，可以看出TensorRT版本需要的是9.2.0.5

组件	版本
Container OS	Ubuntu 22.04
CUDA	NVIDIA CUDA 12.2.2
cuDNN	8.9.5
Python	3.10.12

在autoDL选择合适的显卡和镜像

需要选择支持cuda12.2的显卡（这个一般由英伟达驱动决定，太老的驱动不支持太高的cuda），或者直接用CPU也可以编译，省钱。
需要选择系统为ubuntu 22.04的镜像
最好python也是3.10
选来选区发现有个miniconda的镜像符合，显卡选择的是A5000（或者用纯CPU），内置了cuda11.8和python3.10。除了这个，暂时没看到其他合适的。

准备环境

建议到~/autodl-tmp路径下面操作，这个地方不会占用系统空间。

安装cuda环境

下载cuda 12.2.2

bash

wget https://developer.download.nvidia.com/compute/cuda/12.2.2/local_installers/cuda_12.2.2_535.104.05_linux.run

给予可执行权限

bash

chmod +x cuda_12.2.2_535.104.05_linux.run

正式执行

bash

./cuda_12.2.2_535.104.05_linux.run

输入accept，然后回车
按空格勾选或者取消勾选和上下方向键上下移动，只勾选CUDA Toolkit 12.2，其他均取消勾选，然后选中Install即可。

bash

CUDA Installer                                                               │
│ - [ ] Driver                                                                 │
│      [ ] 535.86.10                                                           │
│ + [X] CUDA Toolkit 12.2                                                      │
│   [ ] CUDA Demo Suite 12.2                                                   │
│   [ ] CUDA Documentation 12.2                                                │
│ - [ ] Kernel Objects                                                         │
│      [ ] nvidia-fs                                                           │
│   Options                                                                    │
│   Install

输入Yes确定。
再次确认cuda版本，看输出已经是cuda12.2.140了

bash

nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Tue_Aug_15_22:02:13_PDT_2023
Cuda compilation tools, release 12.2, V12.2.140
Build cuda_12.2.r12.2/compiler.33191640_0

删除安装包（可选）

bash

rm cuda_12.2.2_535.104.05_linux.run

删除已存在的cuda11.8（可选）

bash

rm -rf /usr/local/cuda-11.8
rm -rf /usr/local/cuda-11

安装cudnn

去英伟达官网下载cudnn，需要登陆英伟达账号（邮箱/微信登陆，没有就注册一个）。官网链接：https://developer.nvidia.com/cudnn
由于系统也内置了cudnn，并且看样子是通过cudnn安装的，所以选择安装包的时候也选择Deb格式,用来覆盖旧版安装会更好。
点击Local Installer for Ubuntu22.04 x86_64 (Deb)下载cudnn 8.9.5 for cuda 12.x & ubuntu 22.04 deb格式的cudnn。
安装cudnn

bash

dpkg -i cudnn-local-repo-ubuntu2204-8.9.5.30_1.0-1_amd64.deb
cd /var/cudnn-local-repo-ubuntu2204-8.9.5.30/
dpkg -i *.deb

删除安装包

bash

rm -rf /var/cudnn-local-repo-ubuntu2204-8.9.5.30/

验证cudnn版本

bash

$ldconfig -v | grep cudnn
libcudnn.so.8 -> libcudnn.so.8.9.5
libcudnn_cnn_infer.so.8 -> libcudnn_cnn_infer.so.8.9.5
libcudnn_adv_train.so.8 -> libcudnn_adv_train.so.8.9.5
libcudnn_ops_train.so.8 -> libcudnn_ops_train.so.8.9.5
libcudnn_ops_infer.so.8 -> libcudnn_ops_infer.so.8.9.5
libcudnn_cnn_train.so.8 -> libcudnn_cnn_train.so.8.9.5
libcudnn_adv_infer.so.8 -> libcudnn_adv_infer.so.8.9.5

安装TensorRT

参考该文档

设置环境变量

bash

export TRT_VER="9.2.0.5"
export CUDA_VER="12.2"
export CUDNN_VER="8.9.4.25-1+cuda12.2"
export NCCL_VER="2.18.3-1+cuda12.2"
export CUBLAS_VER="12.2.5.6-1"
export ARCH="x86_64"
export TRT_CUDA_VERSION="12.2"
export OS="linux"

export RELEASE_URL_TRT=https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/9.2.0/tensorrt-${TRT_VER}.${OS}.${ARCH}-gnu.cuda-${TRT_CUDA_VERSION}.tar.gz;

下载TensorRT并且解压TensorRT

bash

wget ${RELEASE_URL_TRT} -O /tmp/TensorRT.tar
tar -xf /tmp/TensorRT.tar -C /usr/local/

移动解压后的文件到/usr/local/tensorrt/目录，和官方容器保持一致，然后删除tensorrt的压缩包

bash

mv /usr/local/TensorRT-${TRT_VER} /usr/local/tensorrt
rm -rf /tmp/TensorRT.tar

配置ldconfig,然后lib可以被系统找到

bash

echo "/usr/local/tensorrt/lib" > /etc/ld.so.conf.d/tensorrt.conf

确定tensorrt版本是9.1.0

bash

ldconfig -v | grep nvinfer
libnvinfer_vc_plugin.so.9 -> libnvinfer_vc_plugin.so.9.2.0
libnvinfer_dispatch.so.9 -> libnvinfer_dispatch.so.9.2.0
libnvinfer.so.9 -> libnvinfer.so.9.2.0
libnvinfer_plugin.so.9 -> libnvinfer_plugin.so.9.2.0
libnvinfer_lean.so.9 -> libnvinfer_lean.so.9.2.0

设置环境

bash

echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/tensorrt/lib:/usr/local/cuda/lib64/" >> /etc/profile

echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/tensorrt/lib:/usr/local/cuda/lib64/" >> ~/.bashrc

source /etc/profile
source ~/.bashrc
# 测试lib库能否找到
echo ${LD_LIBRARY_PATH}
# 输出结果
# /usr/local/tensorrt/lib:/usr/local/cuda/lib64/

安装并测试mpi4py

安装

bash

conda install mpi4py openmpi

测试

bash

$python3
Python 3.10.8 (main, Nov 24 2022, 14:13:03) [GCC 11.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from mpi4py import MPI
>>> MPI.COMM_WORLD.Get_rank()
0
>>> MPI.COMM_WORLD.Get_size()
1

编译TensorRT-LLM

这里建议还是在/root/autodl-tmp目录编译

下载TensorRT-LLM源码，目前rel为最新稳定分支（如果连不上github，可以本地下载好后，将压缩包放网盘再传上去，然后再解压；也可以用AutoDL自带的镜像加速功能，链接）

bash

cd /root/autodl-tmp
git clone https://github.com/NVIDIA/TensorRT-LLM.git -b rel
# 或者用第三方加速源下载
git clone https://ghproxy.net/https://github.com/NVIDIA/TensorRT-LLM.git -b  rel

bash

cd /root/autodl-tmp/TensorRT-LLM/docker/common

安装基础依赖

bash

apt update
sudo apt-get install -y --no-install-recommends ccache wget gdb git-lfs libffi-dev libopenmpi-dev

安装tensorrt-python，选择和python3.10有关的whl文件。

bash

cd /usr/local/tensorrt/python/
pip install tensorrt-9.2.0.post12.dev5-cp310-none-linux_x86_64.whl
pip install tensorrt_dispatch-9.2.0.post12.dev5-cp310-none-linux_x86_64.whl
pip install tensorrt_lean-9.2.0.post12.dev5-cp310-none-linux_x86_64.whl

测试python tensorrt版本，输出9.2.0.post12.dev5则表示正常

bash

python3 -c "import tensorrt as trt; print(trt.__version__)"

安装polygraphy

bash

cd /root/autodl-tmp/TensorRT-LLM/docker/common/
chmod +x install_polygraphy.sh
./install_polygraphy.sh

一键编译whl文件，编译过程中发现/tmp目录大量占用空间，而AutoDL附赠的30G系统空间不太够用，所以需要指定一个大点的路径来充当tmp路径。

设置数据盘为tmp路径，这个有免费50G空间

bash

cd /root/autodl-tmp/TensorRT-LLM
export TMPDIR=/root/autodl-tmp

拉取第三方子项目代码，如果拉取不了，则按下面的教程修改仓库地址。

bash

git submodule update --init --recursive

修改一下.git/config，加上代理源，否则可能会卡住。

bash

vim .git/config

修改前

bash

[core]
        repositoryformatversion = 0
        filemode = true
        bare = false
        logallrefupdates = true
[remote "origin"]
        url = https://ghproxy.net/https://github.com/NVIDIA/TensorRT-LLM.git
        fetch = +refs/heads/*:refs/remotes/origin/*
[branch "rel"]
        remote = origin
        merge = refs/heads/rel
[submodule "3rdparty/NVTX"]
        active = true
        url = https://github.com/NVIDIA/NVTX.git
[submodule "3rdparty/cutlass"]
        active = true
        url = https://github.com/NVIDIA/cutlass.git
[submodule "3rdparty/cxxopts"]
        active = true
        url = https://github.com/jarro2783/cxxopts
[submodule "3rdparty/json"]
        active = true
        url = https://github.com/nlohmann/json.git

修改后

bash

[core]
        repositoryformatversion = 0
        filemode = true
        bare = false
        logallrefupdates = true
[remote "origin"]
        url = https://ghproxy.net/https://github.com/NVIDIA/TensorRT-LLM.git
        fetch = +refs/heads/*:refs/remotes/origin/*
[branch "rel"]
        remote = origin
        merge = refs/heads/rel
[submodule "3rdparty/NVTX"]
        active = true
        url = https://ghproxy.net/https://github.com/NVIDIA/NVTX.git
[submodule "3rdparty/cutlass"]
        active = true
        url = https://ghproxy.net/https://github.com/NVIDIA/cutlass.git
[submodule "3rdparty/cxxopts"]
        active = true
        url = https://ghproxy.net/https://github.com/jarro2783/cxxopts
[submodule "3rdparty/json"]
        active = true
        url = https://ghproxy.net/https://github.com/nlohmann/json.git

再次手动拉取第三方子仓库源码

bash

git submodule update --init --recursive

手动拉取lfs大文件

bash

git lfs install && git lfs pull

一键编译，大概需要1个小时左右

bash

export BUILD_WHEEL_ARGS="--clean --trt_root /usr/local/tensorrt"

python3 scripts/build_wheel.py ${BUILD_WHEEL_ARGS}

安装whl文件

bash

cd /root/autodl-tmp/TensorRT-LLM/build
pip install tensorrt_llm-0.6.1-py3-none-any.whl

安装后清理一下cache。

bash

# 顺便清理cache
rm ~/.cache/ -rf

测试安装后tensorrt-llm版本

bash

pip list | grep tensorrt-llm

# 输出
# tensorrt-llm             0.6.1

测试tensorrt-llm

bash

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

测试TensorRT-LLM是否可用

将刚刚下载的仓库再下载一份到/root路径，防止后续有发生更新，导致你编译的和它不一样而报错，此时可以不用管第三方仓库和lfs大文件了，因为这个只有编译的时候采用

bash

cd /root
git clone https://ghproxy.net/https://github.com/NVIDIA/TensorRT-LLM.git -b rel

测试tensorrt-llm中的attention，观察tensort-llm可否正常使用（这个操作可能需要显卡）

bash

cd /root/TensorRT-LLM/tests/attention
pip install parameterized
python3 test_bert_attention.py

返回下面这个说明成功了，大概需要5-10分钟才能测试完成，请耐心等待。

bash

----------------------------------------------------------------------
Ran 49 tests in 255.221s

OK

最后，去容器实例关机一下，然后将其保存为镜像，镜像名称随意。

其他教程

更多TensorRT-LLM配套教程可以查看该仓库https://github.com/Tlntin/Qwen-TensorRT-LLM，觉得有用的帮忙点个star哈

背景 ​

准备工作 ​

在autoDL选择合适的显卡和镜像 ​

准备环境 ​

安装cuda环境 ​

安装cudnn ​

安装TensorRT ​

安装并测试mpi4py ​

编译TensorRT-LLM ​

测试TensorRT-LLM是否可用 ​

其他教程 ​

背景

准备工作

在autoDL选择合适的显卡和镜像

准备环境

安装cuda环境

安装cudnn

安装TensorRT

安装并测试mpi4py

编译TensorRT-LLM

测试TensorRT-LLM是否可用

其他教程