fix: build config and depenency

This commit is contained in:
2024-09-11 11:23:42 -07:00
parent 2c69d563be
commit bb2e6c5dc6
2 changed files with 31 additions and 6 deletions
+27 -3
View File
@@ -2,7 +2,9 @@ ARG ROCM_VERSION="5.7.3"
FROM wtongze/rocm-base:${ROCM_VERSION} FROM wtongze/rocm-base:${ROCM_VERSION}
ARG PYTORCH_BRANCH="v2.4.1" ARG ROCM_VERSION="5.7.3"
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"
@@ -27,11 +29,33 @@ EOF
RUN python${PYTHON_VERSION} -m pip install typing-extensions wheel RUN python${PYTHON_VERSION} -m pip install typing-extensions wheel
# For ROCm 4.5+
RUN <<EOF
sudo apt-get update
sudo apt-get install -y rocm-ml-sdk rocm-hip-sdk
sudo apt-get clean
sudo rm -rf /var/lib/apt/lists/*
EOF
RUN <<EOF
sudo apt-get update
sudo apt-get install -y ninja-build
sudo apt-get clean
sudo rm -rf /var/lib/apt/lists/*
EOF
ENV ROCM_PATH=/opt/rocm-${ROCM_VERSION}
ENV PYTORCH_ROCM_ARCH=${HIP_ARCH} ENV PYTORCH_ROCM_ARCH=${HIP_ARCH}
RUN python${PYTHON_VERSION} tools/amd_build/build_amd.py && USE_ROCM=1 python${PYTHON_VERSION} setup.py bdist_wheel ENV PYTORCH_BUILD_VERSION=${PYTORCH_BUILD_VERSION}
ENV PYTORCH_BUILD_NUMBER=0
RUN python${PYTHON_VERSION} tools/amd_build/build_amd.py && \
USE_ROCM=ON USE_FLASH_ATTENTION=OFF USE_MEM_EFF_ATTENTION=OFF python${PYTHON_VERSION} setup.py bdist_wheel
# https://github.com/pytorch/pytorch/issues/125230
# https://github.com/AngryLoki/pytorch/commit/38d6f115c65b67946de6765165ece0288c8416d4
WORKDIR /home/rocm/pytorch/dist WORKDIR /home/rocm/pytorch/dist
RUN cp (ls -l *.whl | awk '{ print $9 }' | tail -n1 | tr -d '\n') /home/rocm/pytorch-rocm.whl RUN cp $(ls -l *.whl | awk '{ print $9 }' | tail -n1 | tr -d '\n') /home/rocm/pytorch-rocm.whl
WORKDIR /home/rocm WORKDIR /home/rocm
CMD [ "bash", "-l" ] CMD [ "bash", "-l" ]
+4 -3
View File
@@ -2,6 +2,7 @@ ARG ROCM_VERSION="5.7.3"
FROM wtongze/rocm-base:${ROCM_VERSION} FROM wtongze/rocm-base:${ROCM_VERSION}
ARG ROCM_VERSION="5.7.3"
ARG TENSORFLOW_BRANCH="r2.12-rocm-enhanced" ARG TENSORFLOW_BRANCH="r2.12-rocm-enhanced"
ARG PYTHON_VERSION="3.11" ARG PYTHON_VERSION="3.11"
ARG NUMPY_VERSION="1.24.3" ARG NUMPY_VERSION="1.24.3"
@@ -47,11 +48,11 @@ RUN sed -i -e "s/python3/python${PYTHON_VERSION}/g" -e "s/pip3/python${PYTHON_VE
RUN sudo ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python RUN sudo ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python
ENV ROCM_PATH=/opt/rocm-${ROCM_VERSION} ENV ROCM_PATH=/opt/rocm-${ROCM_VERSION}
ENV ROCM_TOOLKIT_PATH=/opt/rocm-${ROCM_VERSION} ENV ROCM_TOOLKIT_PATH=${ROCM_PATH}
RUN ./build_rocm_python3 /opt/rocm-${ROCM_VERSION} RUN ./build_rocm_python3 ${ROCM_PATH}
WORKDIR /tmp/tensorflow_pkg WORKDIR /tmp/tensorflow_pkg
RUN cp (ls -l tensorflow*.whl | awk '{ print $9 }' | tail -n1 | tr -d '\n') /home/rocm/tensorflow-rocm.whl RUN cp $(ls -l tensorflow*.whl | awk '{ print $9 }' | tail -n1 | tr -d '\n') /home/rocm/tensorflow-rocm.whl
WORKDIR /home/rocm WORKDIR /home/rocm
CMD [ "/bin/bash", "-l" ] CMD [ "/bin/bash", "-l" ]