Skip to content

Commit

Permalink
Dockerfile.rocm.ubi: get rid of build triton stage
Browse files Browse the repository at this point in the history
this is a torch dependency when installed from the pytorch/rocm6.1
index: https://download.pytorch.org/whl/nightly/rocm6.1
  • Loading branch information
dtrifiro committed Sep 25, 2024
1 parent 69ac6c1 commit 399c114
Showing 1 changed file with 0 additions and 27 deletions.
27 changes: 0 additions & 27 deletions Dockerfile.rocm.ubi
Original file line number Diff line number Diff line change
Expand Up @@ -126,31 +126,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \

##################################################################################################

FROM rocm_devel AS build_triton

# Whether to build triton on rocm
ARG BUILD_TRITON="0"
ARG TRITON_BRANCH="e0fc12c"

# Build triton wheel if `BUILD_TRITON` is set to 1
RUN --mount=type=cache,target=/root/.cache/ccache \
--mount=type=cache,target=/root/.cache/pip \
--mount=type=cache,target=/root/.cache/uv \
if [ "$BUILD_TRITON" = "1" ]; then \
mkdir -p /libs && cd /libs && \
git clone https://github.com/OpenAI/triton.git && \
cd triton && \
git checkout "${TRITON_BRANCH}" && \
cd python && \
uv pip install -v cmake ninja && \
python3 setup.py bdist_wheel --dist-dir=/install; \
else \
mkdir -p /install && \
wget -P /install "https://download.pytorch.org/whl/pytorch_triton_rocm-3.0.0-cp312-cp312-linux_x86_64.whl#sha256=650ce7a6624ecbbe90c30fccb472a57772c08bc6f12312246f7af584f275b7ca"; \
fi

##################################################################################################

FROM rocm_devel AS build_vllm
ARG PYTORCH_ROCM_ARCH
ARG MAX_JOBS
Expand Down Expand Up @@ -188,7 +163,6 @@ RUN microdnf install -y --setopt=install_weak_deps=0 --nodocs gcc && \
microdnf clean all

RUN --mount=type=bind,from=build_amdsmi,src=/install,target=/install/amdsmi/ \
--mount=type=bind,from=build_triton,src=/install,target=/install/triton \
--mount=type=bind,from=build_flashattention,src=/install,target=/install/flashattention \
--mount=type=bind,from=build_vllm,src=/workspace/dist,target=/install/vllm/ \
--mount=type=cache,target=/root/.cache/pip \
Expand All @@ -197,7 +171,6 @@ RUN --mount=type=bind,from=build_amdsmi,src=/install,target=/install/amdsmi/ \
--index-strategy=unsafe-best-match \
--extra-index-url "https://download.pytorch.org/whl/nightly/rocm6.1" \
/install/amdsmi/*.whl\
/install/triton/*.whl\
/install/flashattention/*.whl\
/install/vllm/*.whl

Expand Down

0 comments on commit 399c114

Please sign in to comment.