From 432ca1ecd0f18d6c8fdfe8467c7c667dbfa627ce Mon Sep 17 00:00:00 2001 From: Gregory Shtrasberg Date: Fri, 6 Dec 2024 15:30:03 -0600 Subject: [PATCH] Adjusted Dockerfile.base to the Navi architecture --- Dockerfile.base | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/Dockerfile.base b/Dockerfile.base index 1a87eb91b04b4..d9449843f1bfd 100644 --- a/Dockerfile.base +++ b/Dockerfile.base @@ -4,7 +4,7 @@ FROM ${BASE_IMAGE} AS base ENV PATH=/opt/rocm/llvm/bin:$PATH ENV ROCM_PATH=/opt/rocm ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib: -ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942 +ARG PYTORCH_ROCM_ARCH=gfx1100;gfx1101;gfx1200;gfx1201 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} ARG PYTHON_VERSION=3.12 @@ -29,7 +29,7 @@ RUN apt-get update -y \ RUN pip install -U packaging cmake ninja wheel setuptools Cython FROM base AS build_hipblaslt -ARG HIPBLASLT_BRANCH="507a649" +ARG HIPBLASLT_BRANCH="d43d84a" # Set to "--legacy_hipblas_direct" for ROCm<=6.2 ARG LEGACY_HIPBLASLT_OPTION= RUN git clone https://github.com/ROCm/hipBLASLt @@ -69,8 +69,6 @@ ARG PYTORCH_BRANCH="8bc4033" ARG PYTORCH_VISION_BRANCH="v0.19.1" ARG PYTORCH_REPO="https://github.com/pytorch/pytorch.git" ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git" -ARG FA_BRANCH="c555642" -ARG FA_REPO="https://github.com/ROCm/flash-attention.git" RUN git clone ${PYTORCH_REPO} pytorch RUN cd pytorch && git checkout ${PYTORCH_BRANCH} && \ pip install -r requirements.txt && git submodule update --init --recursive \ @@ -81,14 +79,8 @@ RUN git clone ${PYTORCH_VISION_REPO} vision RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \ && python3 setup.py bdist_wheel --dist-dir=dist \ && pip install dist/*.whl -RUN git clone ${FA_REPO} -RUN cd flash-attention \ - && git checkout ${FA_BRANCH} \ - && git submodule update --init \ - && MAX_JOBS=64 GPU_ARCHS=${PYTORCH_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \ - && cp /app/vision/dist/*.whl /app/install \ - && cp /app/flash-attention/dist/*.whl /app/install + && cp /app/vision/dist/*.whl /app/install FROM base AS final RUN --mount=type=bind,from=build_hipblaslt,src=/app/install/,target=/install \