在Jetson上安装vLLM

Jetson Orin AGX 下编译相关 vLLM 软件

编译 PyTorch

如果想用 vllm,官方提供的 PyTorch 二进制安装包是不行的,因为它需要用到 distributed 以及 NCCL 的支持,官方提供的包中是不包含这些特性的,甚至一些民间爱好者提供的安装包也只支持了 distributed 而未支持 NCCL,所以需要自行安装。

下载必要的工具

需要下载安装必要的工具,因为 Jetson Linux 是 Ubuntu 系统,所以直接使用 apt-get 即可。

sudo apt-get install python3-pip cmake libopenblas-dev libopenmpi-dev

注:Jetson Linux 自带开发相关的软件包,比如 git, gcc 等,如果是在普通的 PC 版 Linux 下,也可以编译,但需要执行多安装 git build-essential这两个 package。

克隆 PyTorch 代码

系统自带 Git,可以直接用以下命令克隆 PyTorch 代码

git clone --recursive http://github.com/pytorch/pytorch

由于默认克隆的是当前开发版本,所以通常需要指定一个版本,如 v2.3.1,则执行

git checkout v2.3.1

切换到该版本,也可直接在克隆项目时添加 ---branch v2.3.1参数指定。

编译 PyTorch 代码

首先确保必须的 Python 包都有安装,并且为了不影响其他环境,先要利用 env 来创建一个虚拟环境。

$ pip3 install -r requirements.txt
$ pip3 install scikit-build
$ pip3 install ninja

编译代码需要设定一些环境变量,所以最方便的做法是编写一个编译脚本build.sh

export USE_NCCL=1
# export USE_SYSTEM_NCCL=0        # if you do not use system nccl, nvidia-smi may do not work
export USE_DISTRIBUTED=1                    # set 0 if you want to disable OpenMPI backend
export USE_QNNPACK=0
export USE_PYTORCH_QNNPACK=0
export TORCH_CUDA_ARCH_LIST="5.2;6.1;7.0;7.2;7.5;8.0;8.6;8.7"
export PYTORCH_BUILD_VERSION=2.3.1  # without the leading 'v', e.g. 2.3.0 for PyTorch v2.3.0
export PYTORCH_BUILD_NUMBER=1

python3 setup.py install

编译 vllm

vllm 需要先编译以下几个依赖库:

  • triton
  • vllm-flash-attn 所以我们将一一介绍。

Triton

triton 实际可以直接编译,但它需要下载大量文件,而这些文件是在境外服务器的,所以,最好是设置好代理服务器,比如

export http_proxy=http://localhost:7890
export https_proxy=http://localhost:7890

然后需要将项目 clone 到本地,因为 pypi.org 未提供代码,也只提供了常用平台的二进制安装包,所以直接 pip 是会报错说找不到相应的安装包的。可以按照以下命令

git clone https://github.com/triton-lang/triton.git;
cd triton;

pip install ninja cmake wheel; # build-time dependencies
pip install -e python

来进行安装。

注意,由于要下载好几个 G 的东西,所以非常耗时,需要做好心理准备。

vllm-flash-attn

虽然实际上,没有vllm-flash-attn,vllm 也会使用内置的 xformers 支持来进行加速。但实际使用下来,在 Orin AGX 下,xformers 是无法正常运行的。所以必须安装vll-flash-attn

但该库不能直接从 pypi.org 上下载编译,需要修改部分设置,所以先需要 clone 下来。

git clone --recurse-submodules https://github.com/vllm-project/flash-attention.git

然后修改 setup.py 中的两行:

一个是if bare_metal_version >= Version("11.8"):下,将 90 改成 87。即改为:

cc_flag.append("arch=compute_87,code=sm_87")

因为现在的 vllm 是需要的 PyTorch 版本是 v2.3.0 的,而最新的已经需要 v2.3.1,所以需要修改后面PYTORCH_VERSION的值从 2.3.1 改成 2.3.0。

修改完之后执行

python setup.py install

进行安装。如果发生内存不足编译失败,可以加MAX_JOBS来限制最大的并行编译量,比如

MAX_JOBS=4 python setup.py install

来限制最大 4 进程编译。

编译 vllm

先将 vllm 的代码 clone 到本地

git clone --recurse-submodules https://github.com/vllm-project/vllm.git

修改 vllm 的 CMakeLists.txt,将其中的CUDA_SUPPORTED_ARCHS中,添加 8.7,可以改成如下:

set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0")

然后可以进行编译

python setup.py install

生成 wheel 安装包

注意,以上的示例都是最后执行python setup.py install直接进行安装的。但如果想生成 wheel 安装包供以后再次安装,或分发给其他用户,那么可以将其改为:

python setup.py bdist_wheel

这样就会在 dist 目录下生成二进制安装包,然后进入该目录,执行

pip install xxx.whl # xxx为该安装包名称

最后修改于 2024-07-30