diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index d923693d11..0bba4be5bb 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -1,11 +1,11 @@ ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete -ARG TRITON_BRANCH="0f380657" +ARG TRITON_BRANCH="0263a6a" # release/internal/3.7.x as of 08/02 ARG TRITON_REPO="https://github.com/ROCm/triton.git" -ARG PYTORCH_BRANCH="d0c8b1f3" # release/2.11 as of 6/09 +ARG PYTORCH_BRANCH="6bbd260" # release/2.12 as of 08/02 ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git" -ARG PYTORCH_VISION_BRANCH="v0.24.1" +ARG PYTORCH_VISION_BRANCH="v0.27.1" ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git" -ARG PYTORCH_AUDIO_BRANCH="v2.9.0" +ARG PYTORCH_AUDIO_BRANCH="v2.11.0" ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git" ARG FA_BRANCH="0e60e394" ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git" @@ -42,16 +42,16 @@ RUN mkdir -p /app WORKDIR /app ENV DEBIAN_FRONTEND=noninteractive -# Install Python and other dependencies +# Install Python and other dependencies (plus gcc-13, selected later in build_pytorch). RUN apt-get update -y \ && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev pkg-config \ && for i in 1 2 3; do \ - add-apt-repository -y ppa:deadsnakes/ppa && break || \ + add-apt-repository -y ppa:deadsnakes/ppa && add-apt-repository -y ppa:ubuntu-toolchain-r/test && break || \ { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ done \ && apt-get update -y \ && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ - python${PYTHON_VERSION}-lib2to3 python-is-python3 \ + python${PYTHON_VERSION}-lib2to3 python-is-python3 gcc-13 g++-13 \ && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ @@ -112,12 +112,9 @@ FROM base AS build_triton ARG TRITON_BRANCH ARG TRITON_REPO RUN git clone ${TRITON_REPO} -# Cherry picking the following -# https://github.com/triton-lang/triton/pull/8991 RUN cd triton \ && git checkout ${TRITON_BRANCH} \ && git config --global user.email "you@example.com" && git config --global user.name "Your Name" \ - && git cherry-pick 555d04f \ && if [ ! -f setup.py ]; then cd python; fi \ && python3 setup.py bdist_wheel --dist-dir=dist \ && mkdir -p /app/install && cp dist/*.whl /app/install @@ -145,6 +142,8 @@ ARG PYTORCH_AUDIO_REPO ARG USE_SCCACHE RUN apt-get update && apt-get install -y pkg-config liblzma-dev +# torch 2.12 emits float->__bf16 libcalls (__truncsfbf2) absent from gcc-11 libgcc; build with gcc-13. +RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-13 130 --slave /usr/bin/g++ g++ /usr/bin/g++-13 RUN git clone ${PYTORCH_REPO} pytorch RUN cd pytorch && git checkout ${PYTORCH_BRANCH} RUN cd pytorch \ @@ -316,4 +315,4 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \ && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \ - && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt + && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt \ No newline at end of file diff --git a/vllm/compilation/caching.py b/vllm/compilation/caching.py index 62da2d9de3..94456fc182 100644 --- a/vllm/compilation/caching.py +++ b/vllm/compilation/caching.py @@ -22,6 +22,7 @@ from vllm.compilation.counter import compilation_counter from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.utils import hash_factors from vllm.logger import init_logger +from vllm.platforms import current_platform from vllm.utils.hashing import safe_hash try: @@ -130,9 +131,16 @@ class StandaloneCompiledArtifacts: compilation_counter.num_compiled_artifacts_loaded += 1 return AOTCompiledArtifact.deserialize(entry) - with concurrent.futures.ThreadPoolExecutor() as executor: - entries = list(self.submodule_bytes_store.values()) - loaded_entries = list(executor.map(_load_entry, entries)) + entries = list(self.submodule_bytes_store.values()) + if current_platform.is_rocm(): + # Deserializing an artifact loads its compiled Triton kernels, which + # on ROCm ends up in hipModuleLoad. Loading modules concurrently + # deadlocks on the dynamic linker lock, since rocprofiler-sdk calls + # dl_iterate_phdr from inside the load, so load one at a time. + loaded_entries = [_load_entry(entry) for entry in entries] + else: + with concurrent.futures.ThreadPoolExecutor() as executor: + loaded_entries = list(executor.map(_load_entry, entries)) for i, k in enumerate(self.submodule_bytes_store.keys()): self.loaded_submodule_store[k] = loaded_entries[i]