使用 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从源码构建

  2. 从源码构建

选项 2:从源代码构建#

  1. 安装先决条件(如果你已经在包含以下安装的环境/docker中,请跳过):

要安装 PyTorch,你可以从一个全新的 Docker 镜像开始,例如 rocm/pytorch:rocm6.2_ubuntu20.04_py3.9_pytorch_release_2.3.0rocm/pytorch-nightly

或者,您可以使用 PyTorch 轮子安装 PyTorch。您可以在 PyTorch 入门指南 中查看 PyTorch 安装指南。

  1. 安装 Triton flash attention for ROCm

按照 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 错误是间歇性的。

  1. 可选地,如果你选择使用 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

  1. 构建 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 性能优化