Сборка последней версии vLLM из исходников под CUDA 11.8 и Torch 2.5.1

vllm-logo

У нас есть проблема с CUDA.

У нас 11.8 версия в кластере, и на текущий момент обновить на 12+ не можем, и по этому приходится выкручиваться различными способами, зачастую болезненными.

Итак. Нужен свежий образ vLLM, который будет там инферить модели внутри kServe. Свежак умеет инферить модели Score и Embedding.

У разрабов vLLM на гитхабе поломались их пайплайны создания релизов, и они их чинить пока не хотят, то есть есть нет нормального whl под CUDA 11.8 и нашего петухона. alt text Помимо всего этого есть другая проблема - нам нужен vllm-flash-attention (форк основной репы, который сделали сами разрабы vLLM) для ускорения инференса LLM моделей. У него в требованиях CUDA 12+ и Torch 2.4.0, при этом же vLLM требует Torch 2.5.1. Уже попахивает шизой.

Собственно решение всего такое:

Тырим логику из workflows гитхаба vLLM и собираем образ внутри которого будет собираться наш whl с нашими требованиями. Собирается из исходников с компиляцией всего С++ и CUDA кода. Занимает часа 2-3. Можно отредактировать MAX_JOBS в зависимости от количества ядер на сборочной машине.

FROM nvidia/cuda:11.8.0-base-ubuntu22.04

LABEL version="1.0"
LABEL description="VLLM CUDA 11.8 Torch 2.5.1"

ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && apt-get install -y \
    python3.10 \
    python3.10-dev \
    python3-pip \
    git \
    wget \
    ninja-build \
    build-essential \
    cmake \
    && rm -rf /var/lib/apt/lists/*


RUN wget -nv https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb \
    && dpkg -i cuda-keyring_1.1-1_all.deb \
    && rm cuda-keyring_1.1-1_all.deb \
    && apt-get update \
    && DEBIAN_FRONTEND=noninteractive apt-get install -y cuda-11-8 cuda-nvcc-11-8 cuda-libraries-dev-11-8 \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

ENV PATH=/usr/local/cuda-11.8/bin:${PATH}
ENV LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:${LD_LIBRARY_PATH}
ENV CUDA_HOME=/usr/local/cuda-11.8

RUN python3.10 -m pip install --no-cache-dir \
    numpy \
    pyyaml \
    scipy \
    ipython \
    mkl \
    mkl-include \
    ninja \
    cython \
    typing \
    pandas \
    typing-extensions \
    dataclasses \
    setuptools \
    wheel

RUN python3.10 -m pip install --no-cache-dir \
    torch==2.5.1+cu118 \
    torchvision==0.20.1+cu118 \
    --extra-index-url https://download.pytorch.org/whl/cu118

RUN git clone https://github.com/vllm-project/vllm.git /vllm
WORKDIR /vllm

RUN python3.10 -m pip install --no-cache-dir -r requirements-build.txt -r requirements-cuda.txt

ENV MAX_JOBS=4
ENV TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
ENV VLLM_FA_CMAKE_GPU_ARCHES="80-real;90-real"

RUN python3.10 setup.py bdist_wheel --dist-dir=dist

RUN mkdir /wheels && cp dist/*.whl /wheels/

RUN python3.10 -m pip install /wheels/*.whl

RUN python3.10 -c "import vllm; print(vllm.__version__)"

############################# WARNING #########################################
# ПОСЛЕ СБОРКИ ЗАБРАТЬ WHL ПАКЕТ ИЗ ОБРАЗА
###############################################################################

Забираем whl из образа.

Собираем образ с vllm-flash-attention. Предварительно нужно стянуть репозиторий vllm-flash-attention и в setup.py версию торча поменять на совместимую с vLLM - в моем случае это 2.5.1

FROM nvidia/cuda:11.8.0-base-ubuntu22.04 AS vllm-base
LABEL version="1.0"
LABEL description="VLLM-Flash-Attn CUDA 11.8 Torch 2.5.1"
ENV CUDA_HOME=/usr/local/cuda-11.8
ENV PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cu118

RUN apt update -y \
    && apt upgrade -y \
    && apt install -y \
        cuda-nvcc-11-8 \
        python3.10 \
        python3.10-dev \
        python3-pip \
        python3.10-venv \
        git \
    && rm -rf /var/lib/apt/lists/* \
    && ln -sf /usr/bin/python3.10 /usr/bin/python3 \
    && ln -sf /usr/bin/python3.10 /usr/bin/python

RUN python3 --version
RUN --mount=type=cache,target=/root/.cache/pip pip install torch==2.5.1 torchvision==0.20.1 transformers==4.43.2

RUN git clone https://github.com/vllm-project/vllm.git /vllm
WORKDIR /vllm
ENV VLLM_USE_PRECOMPILED=1
RUN --mount=type=cache,target=/root/.cache/pip pip install --editable .

FROM vllm-base AS vllm-flash-attn-build
ARG max_jobs=4
ENV MAX_JOBS=${max_jobs}

############################# WARNING #########################################
# ВНИМАНИЕ! НЕОБХОДИМО СКАЧАТЬ РЕПОЗИТОРИЙ VLLM-FLASH-ATTENTION ЛОКАЛЬНО И В ФАЙЛЕ SETUP.PY ИЗМЕНИТЬ ВЕРСИЮ TORCH НА 2.5.1 (ИЛИ АКТУАЛЬНУЮ)
# ВНИМАНИЕ! НЕОБХОДИМО СКАЧАТЬ РЕПОЗИТОРИЙ VLLM-FLASH-ATTENTION ЛОКАЛЬНО И В ФАЙЛЕ SETUP.PY ИЗМЕНИТЬ ВЕРСИЮ TORCH НА 2.5.1 (ИЛИ АКТУАЛЬНУЮ)
# ВНИМАНИЕ! НЕОБХОДИМО СКАЧАТЬ РЕПОЗИТОРИЙ VLLM-FLASH-ATTENTION ЛОКАЛЬНО И В ФАЙЛЕ SETUP.PY ИЗМЕНИТЬ ВЕРСИЮ TORCH НА 2.5.1 (ИЛИ АКТУАЛЬНУЮ)
###############################################################################

COPY flash-attention ./flash-attention
RUN apt update -y \
    && apt install -y git cuda-libraries-dev-11-8 \
    && pip install packaging ninja
WORKDIR flash-attention
RUN python3 setup.py bdist_wheel --dist-dir=/vllm-flash-attn-dist

FROM vllm-base AS vllm-runtime
COPY --from=vllm-flash-attn-build /vllm-flash-attn-dist /vllm-flash-attn-dist
RUN pip install /vllm-flash-attn-dist/*.whl

############################# WARNING #########################################
# ПОСЛЕ СБОРКИ ЗАБРАТЬ WHL ПАКЕТ ИЗ ОБРАЗА
###############################################################################

Не забываем забрать whl из образа.

Собираем итоговый образ с vLLM и vllm-flash-attention под CUDA 11.8.

FROM nvidia/cuda:11.8.0-base-ubuntu22.04 AS vllm-base

LABEL version="1.1"
LABEL description="VLLM runtime with Flash Attention CUDA 11.8"


ENV CUDA_HOME=/usr/local/cuda-11.8
ENV PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cu118

RUN apt update -y \
    && apt upgrade -y \
    && apt install -y \
        cuda-nvcc-11-8 \
        python3.10 \
        python3.10-dev \
        python3-pip \
        python3.10-venv \
    && rm -rf /var/lib/apt/lists/* \
    && ln -sf /usr/bin/python3.10 /usr/bin/python3 \
    && ln -sf /usr/bin/python3.10 /usr/bin/python

RUN python3 --version

RUN --mount=type=cache,target=/root/.cache/pip pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu118
RUN --mount=type=cache,target=/root/.cache/pip pip install xformers --index-url https://download.pytorch.org/whl/cu118
RUN --mount=type=cache,target=/root/.cache/pip pip install packaging
RUN apt update && apt install -y git wget cuda-libraries-dev-11-8 && \
    wget https://github.com/Kitware/CMake/releases/download/v3.26.0/cmake-3.26.0-linux-x86_64.sh -q && \
    chmod +x cmake-3.26.0-linux-x86_64.sh && \
    ./cmake-3.26.0-linux-x86_64.sh --skip-license --prefix=/usr && \
    rm cmake-3.26.0-linux-x86_64.sh
RUN --mount=type=cache,target=/root/.cache/pip pip install --upgrade setuptools>=61
RUN --mount=type=cache,target=/root/.cache/pip pip install setuptools_scm



############################# WARNING #########################################
# WHL ПАКЕТЫ НЕОБХОДИМО СОБРАТЬ ПРЕДВАРИТЕЛЬНО. DOCKERFILE ДЛЯ СБОРКИ WHL ЛЕЖАТ РЯДОМ 
###############################################################################

COPY vllm-0.6.4.post2.dev395+g02222a02.cu118-cp310-cp310-linux_x86_64.whl /vllm/vllm-0.6.4.post2.dev395+g02222a02.cu118-cp310-cp310-linux_x86_64.whl
COPY vllm_flash_attn-2.6.1+cu118-cp310-cp310-linux_x86_64.whl /vllm/vllm_flash_attn-2.6.1+cu118-cp310-cp310-linux_x86_64.whl

WORKDIR /vllm
RUN pip install vllm-0.6.4.post2.dev395+g02222a02.cu118-cp310-cp310-linux_x86_64.whl
RUN pip install vllm_flash_attn-2.6.1+cu118-cp310-cp310-linux_x86_64.whl

На этом собственно все. Это базовый вариант, без multistage build. Докерфайлы можно оптимизировать, но это уже на ваше усмотрение.