[ROCm]: Bump torch 2.12, triton 3.7, torchaudio, torchvision (#50607)

Signed-off-by: Rohan Potdar <rohan.potdar@amd.com>
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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

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]