使用 ROCm 安装#
vLLM 支持使用 ROCm 6.2 的 AMD GPU。
要求#
操作系统: Linux
Python: 3.9 – 3.12
GPU: MI200s (gfx90a), MI300 (gfx942), Radeon RX 7900 系列 (gfx1100)
ROCm 6.2
注意:PyTorch 2.5+/ROCm6.2 已放弃对 Python 3.8 的支持。
安装选项:
选项 1:使用 Docker 从源代码构建(推荐)#
你可以从源代码构建并安装 vLLM。
首先,从 Dockerfile.rocm 构建一个 Docker 镜像,并从该镜像启动一个 Docker 容器。重要的是,用户使用 buildkit 启动 Docker 构建。用户可以在调用 docker build 命令时将 DOCKER_BUILDKIT=1 作为环境变量,或者用户需要在 docker daemon 配置 /etc/docker/daemon.json 中进行如下设置并重启 daemon:
{
"features": {
"buildkit": true
}
}
Dockerfile.rocm 默认使用 ROCm 6.2,但也支持旧版 vLLM 分支中的 ROCm 5.7、6.0 和 6.1。它提供了灵活性,可以使用以下参数自定义 docker 镜像的构建:
BASE_IMAGE: 指定运行
docker build
时使用的基本镜像,特别是 PyTorch on ROCm 基本镜像。BUILD_FA: 指定是否构建 CK flash-attention。默认值为 1。对于 Radeon RX 7900 系列 (gfx1100),在 flash-attention 支持此目标之前,应将其设置为 0。
FX_GFX_ARCHS: 指定用于构建CK flash-attention的GFX架构,例如,MI200和MI300使用`gfx90a;gfx942`。默认值为`gfx90a;gfx942`
FA_BRANCH: 指定用于在 ROCm 的 flash-attention 仓库 中构建 CK flash-attention 的分支。默认是 ae7928c
BUILD_TRITON: 指定是否构建 triton flash-attention。默认值为 1。
它们的值可以在运行 docker build
时通过 --build-arg
选项传递。
要在 ROCm 6.2 上为 MI200 和 MI300 系列构建 vllm,您可以使用默认设置。
$ DOCKER_BUILDKIT=1 docker build -f Dockerfile.rocm -t vllm-rocm .
要在 Radeon RX7900 系列(gfx1100)上为 ROCm 6.2 构建 vllm,您应如下指定 BUILD_FA
:
$ DOCKER_BUILDKIT=1 docker build --build-arg BUILD_FA="0" -f Dockerfile.rocm -t vllm-rocm .
要运行上述docker镜像 vllm-rocm
,请使用以下命令:
$ docker run -it \
--network=host \
--group-add=video \
--ipc=host \
--cap-add=SYS_PTRACE \
--security-opt seccomp=unconfined \
--device /dev/kfd \
--device /dev/dri \
-v <path/to/model>:/app/model \
vllm-rocm \
bash
其中 <path/to/model> 是模型存储的位置,例如,llama2 或 llama3 模型的权重。
选项 2:从源代码构建#
安装先决条件(如果你已经在包含以下安装的环境/docker中,请跳过):
要安装 PyTorch,你可以从一个全新的 Docker 镜像开始,例如 rocm/pytorch:rocm6.2_ubuntu20.04_py3.9_pytorch_release_2.3.0,rocm/pytorch-nightly。
或者,您可以使用 PyTorch 轮子安装 PyTorch。您可以在 PyTorch 入门指南 中查看 PyTorch 安装指南。
按照 ROCm/triton 的说明安装 ROCm 的 Triton flash attention(默认的 triton-mlir 分支)。
$ python3 -m pip install ninja cmake wheel pybind11 $ pip uninstall -y triton $ git clone https://github.com/OpenAI/triton.git $ cd triton $ git checkout e192dba $ cd python $ pip3 install . $ cd ../..
备注
如果在构建 Triton 期间遇到与下载包相关的 HTTP 问题,请重试,因为 HTTP 错误是间歇性的。
可选地,如果你选择使用 CK flash attention,你可以安装 flash attention for ROCm
按照 ROCm/flash-attention 的说明安装 ROCm 的 flash attention (v2.5.9.post1)。或者,可以在发布中获取用于 vLLM 使用的轮子。
例如,对于 ROCm 6.2,假设你的 gfx 架构是 gfx90a。注意要获取你的 gfx 架构,请运行 rocminfo |grep gfx。
$ git clone https://github.com/ROCm/flash-attention.git $ cd flash-attention $ git checkout 3cea2fb $ git submodule update --init $ GPU_ARCHS="gfx90a" python3 setup.py install $ cd ..
备注
你可能需要将“ninja”版本降级到1.10,因为在编译flash-attention-2时不会使用它(例如 pip install ninja==1.10.2.4)
构建 vLLM。
例如,vLLM 在 ROCM 6.2 上可以通过以下步骤构建:
$ pip install --upgrade pip $ # Install PyTorch $ pip uninstall torch -y $ pip install --no-cache-dir --pre torch==2.6.0.dev20240918 --index-url https://download.pytorch.org/whl/nightly/rocm6.2 $ # Build & install AMD SMI $ pip install /opt/rocm/share/amd_smi $ # Install dependencies $ pip install --upgrade numba scipy huggingface-hub[cli] $ pip install "numpy<2" $ pip install -r requirements-rocm.txt $ # Build vLLM for MI210/MI250/MI300. $ export PYTORCH_ROCM_ARCH="gfx90a;gfx942" $ python3 setup.py develop
这可能需要5-10分钟。目前,
pip install .
不适用于ROCm安装。
小技巧
Triton flash attention 默认情况下被使用。为了基准测试的目的,建议在收集性能数据之前运行一个预热步骤。
Triton flash attention 目前不支持滑动窗口注意力。如果使用半精度,请使用 CK flash-attention 以支持滑动窗口。
要使用 CK flash-attention 或 PyTorch naive attention,请使用此标志
export VLLM_USE_TRITON_FLASH_ATTN=0
来关闭 triton flash attention。理想情况下,PyTorch 的 ROCm 版本应与 ROCm 驱动程序版本匹配。
小技巧
对于 MI300x (gfx942) 用户,为了达到最佳性能,请参考 MI300x 调优指南 以获取系统和工作流级别的性能优化和调优技巧。对于 vLLM,请参考 vLLM 性能优化。