vllm
/
Dockerfile.rocm
95 строк · 3.1 Кб
1# default base image
2ARG BASE_IMAGE="rocm/pytorch:rocm6.0_ubuntu20.04_py3.9_pytorch_2.1.1"
3
4FROM $BASE_IMAGE
5
6ARG BASE_IMAGE="rocm/pytorch:rocm6.0_ubuntu20.04_py3.9_pytorch_2.1.1"
7
8RUN echo "Base image is $BASE_IMAGE"
9
10# BASE_IMAGE for ROCm_5.7: "rocm/pytorch:rocm5.7_ubuntu22.04_py3.10_pytorch_2.0.1"
11# BASE_IMAGE for ROCm_6.0: "rocm/pytorch:rocm6.0_ubuntu20.04_py3.9_pytorch_2.1.1"
12
13
14ARG FA_GFX_ARCHS="gfx90a;gfx942"
15RUN echo "FA_GFX_ARCHS is $FA_GFX_ARCHS"
16
17ARG FA_BRANCH="3d2b6f5"
18RUN echo "FA_BRANCH is $FA_BRANCH"
19
20# whether to build flash-attention
21# if 0, will not build flash attention
22# this is useful for gfx target where flash-attention is not supported
23# In that case, we need to use the python reference attention implementation in vllm
24ARG BUILD_FA="1"
25
26# Install some basic utilities
27RUN apt-get update && apt-get install python3 python3-pip -y
28
29# Install some basic utilities
30RUN apt-get update && apt-get install -y \
31curl \
32ca-certificates \
33sudo \
34git \
35bzip2 \
36libx11-6 \
37build-essential \
38wget \
39unzip \
40nvidia-cuda-toolkit \
41tmux \
42&& rm -rf /var/lib/apt/lists/*
43
44### Mount Point ###
45# When launching the container, mount the code directory to /app
46ARG APP_MOUNT=/app
47VOLUME [ ${APP_MOUNT} ]
48WORKDIR ${APP_MOUNT}
49
50RUN python3 -m pip install --upgrade pip
51RUN python3 -m pip install --no-cache-dir fastapi ninja tokenizers pandas
52
53ENV LLVM_SYMBOLIZER_PATH=/opt/rocm/llvm/bin/llvm-symbolizer
54ENV PATH=$PATH:/opt/rocm/bin:/libtorch/bin:
55ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/rocm/lib/:/libtorch/lib:
56ENV CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:/libtorch/include:/libtorch/include/torch/csrc/api/include/:/opt/rocm/include/:
57
58# Install ROCm flash-attention
59RUN if [ "$BUILD_FA" = "1" ]; then \
60mkdir libs \
61&& cd libs \
62&& git clone https://github.com/ROCm/flash-attention.git \
63&& cd flash-attention \
64&& git checkout ${FA_BRANCH} \
65&& git submodule update --init \
66&& export GPU_ARCHS=${FA_GFX_ARCHS} \
67&& if [ "$BASE_IMAGE" = "rocm/pytorch:rocm5.7_ubuntu22.04_py3.10_pytorch_2.0.1" ]; then \
68patch /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/utils/hipify/hipify_python.py hipify_patch.patch; fi \
69&& python3 setup.py install \
70&& cd ..; \
71fi
72
73COPY ./ /app/vllm
74
75RUN python3 -m pip install --upgrade pip
76RUN python3 -m pip install xformers==0.0.23 --no-deps
77
78# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt.
79# Manually removed it so that later steps of numpy upgrade can continue
80RUN if [ "$BASE_IMAGE" = "rocm/pytorch:rocm6.0_ubuntu20.04_py3.9_pytorch_2.1.1" ]; then \
81rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/; fi
82
83RUN cd /app \
84&& cd vllm \
85&& pip install -U -r requirements-rocm.txt \
86&& if [ "$BUILD_FA" = "1" ]; then \
87bash patch_xformers.rocm.sh; fi \
88&& patch /opt/rocm/include/hip/amd_detail/amd_hip_bf16.h /app/vllm/rocm_patch/rocm_bf16.patch \
89&& python3 setup.py install \
90&& cd ..
91
92RUN python3 -m pip install --upgrade pip
93RUN python3 -m pip install --no-cache-dir ray[all]
94
95CMD ["/bin/bash"]
96