feat: simplify pytorch version

This commit is contained in:
2024-09-12 12:32:37 -07:00
parent 259ec79cf4
commit b55c359194
+5 -3
View File
@@ -4,7 +4,6 @@ FROM wtongze/rocm-base:${ROCM_VERSION}
ARG ROCM_VERSION="5.7.3" ARG ROCM_VERSION="5.7.3"
ARG PYTORCH_BRANCH="v2.2.2" ARG PYTORCH_BRANCH="v2.2.2"
ARG PYTORCH_BUILD_VERSION="2.2.2"
ARG PYTHON_VERSION="3.11" ARG PYTHON_VERSION="3.11"
ARG HIP_ARCH="gfx803;gfx900" ARG HIP_ARCH="gfx803;gfx900"
ARG DEBIAN_FRONTEND="noninteractive" ARG DEBIAN_FRONTEND="noninteractive"
@@ -38,10 +37,13 @@ EOF
ENV ROCM_PATH=/opt/rocm-${ROCM_VERSION} ENV ROCM_PATH=/opt/rocm-${ROCM_VERSION}
ENV PYTORCH_ROCM_ARCH=${HIP_ARCH} ENV PYTORCH_ROCM_ARCH=${HIP_ARCH}
ENV PYTORCH_BUILD_VERSION=${PYTORCH_BUILD_VERSION}
ENV PYTORCH_BUILD_NUMBER=0 ENV PYTORCH_BUILD_NUMBER=0
RUN python${PYTHON_VERSION} tools/amd_build/build_amd.py && \ RUN python${PYTHON_VERSION} tools/amd_build/build_amd.py && \
USE_ROCM=ON USE_NUMPY=1 USE_FLASH_ATTENTION=OFF USE_MEM_EFF_ATTENTION=OFF python${PYTHON_VERSION} setup.py bdist_wheel USE_ROCM=ON \
USE_NUMPY=1 \
USE_FLASH_ATTENTION=OFF \
USE_MEM_EFF_ATTENTION=OFF \
PYTORCH_BUILD_VERSION=$(echo "${PYTORCH_BRANCH}" | tr -d "v") python${PYTHON_VERSION} setup.py bdist_wheel
# https://github.com/pytorch/pytorch/issues/125230 # https://github.com/pytorch/pytorch/issues/125230
# https://github.com/AngryLoki/pytorch/commit/38d6f115c65b67946de6765165ece0288c8416d4 # https://github.com/AngryLoki/pytorch/commit/38d6f115c65b67946de6765165ece0288c8416d4