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