95 lines
2.7 KiB
Docker
95 lines
2.7 KiB
Docker
ARG ROCM_VERSION="5.7.3"
|
|
|
|
FROM wtongze/rocm-base:${ROCM_VERSION}
|
|
|
|
ARG ROCM_VERSION="5.7.3"
|
|
ARG PYTORCH_BRANCH="v2.2.2"
|
|
ARG TORCHVISION_BRANCH="v0.17.2"
|
|
ARG PYTHON_VERSION="3.11"
|
|
ARG NUMPY_VERSION="1.26.4"
|
|
ARG HIP_ARCH="gfx803;gfx900"
|
|
ARG DEBIAN_FRONTEND="noninteractive"
|
|
|
|
# Add GPU target
|
|
RUN echo "${HIP_ARCH}" | sed 's/;/\n/g' | sudo tee /opt/rocm/bin/target.lst
|
|
|
|
WORKDIR /home/rocm
|
|
RUN git clone --depth 1 https://github.com/pytorch/pytorch.git -b ${PYTORCH_BRANCH} pytorch
|
|
WORKDIR /home/rocm/pytorch
|
|
RUN git submodule update --init --recursive
|
|
|
|
RUN <<EOF
|
|
sudo apt-get update
|
|
sudo apt install -y software-properties-common
|
|
sudo add-apt-repository -y ppa:deadsnakes/ppa
|
|
sudo apt-get update
|
|
sudo apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev
|
|
sudo apt-get clean
|
|
sudo rm -rf /var/lib/apt/lists/*
|
|
EOF
|
|
|
|
RUN python${PYTHON_VERSION} -m pip install wheel
|
|
RUN python${PYTHON_VERSION} -m pip install -r requirements.txt
|
|
RUN python${PYTHON_VERSION} -m pip install numpy==${NUMPY_VERSION}
|
|
|
|
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_BUILD_NUMBER=0
|
|
|
|
# Torch
|
|
# > https://github.com/pytorch/pytorch/issues/125230
|
|
# > https://github.com/AngryLoki/pytorch/commit/38d6f115c65b67946de6765165ece0288c8416d4
|
|
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 \
|
|
PYTORCH_BUILD_VERSION=$(echo "${PYTORCH_BRANCH}" | tr -d "v") python${PYTHON_VERSION} setup.py bdist_wheel
|
|
|
|
RUN <<EOF
|
|
sudo apt-get update
|
|
sudo apt-get install -y ffmpeg libjpeg-dev libpng-dev
|
|
sudo apt-get clean
|
|
sudo rm -rf /var/lib/apt/lists/*
|
|
EOF
|
|
|
|
RUN python${PYTHON_VERSION} -m pip install /home/rocm/pytorch/dist/*.whl
|
|
|
|
# Torch Vision
|
|
WORKDIR /home/rocm/
|
|
RUN git clone --depth 1 https://github.com/pytorch/vision.git -b ${TORCHVISION_BRANCH} vision
|
|
|
|
WORKDIR /home/rocm/vision
|
|
RUN USE_ROCM=ON \
|
|
BUILD_VERSION=$(echo "${TORCHVISION_BRANCH}" | tr -d "v") \
|
|
python${PYTHON_VERSION} setup.py bdist_wheel
|
|
|
|
# Torch Audio
|
|
WORKDIR /home/rocm/
|
|
RUN git clone --depth 1 https://github.com/pytorch/audio.git -b ${PYTORCH_BRANCH} audio
|
|
|
|
WORKDIR /home/rocm/audio
|
|
RUN python${PYTHON_VERSION} -m pip install -r requirements.txt
|
|
RUN USE_ROCM=ON \
|
|
BUILD_VERSION=$(echo "${PYTORCH_BRANCH}" | tr -d "v") \
|
|
python${PYTHON_VERSION} setup.py bdist_wheel
|
|
|
|
RUN <<EOF
|
|
mkdir /home/rocm/build
|
|
cp /home/rocm/pytorch/dist/*.whl /home/rocm/build
|
|
cp /home/rocm/vision/dist/*.whl /home/rocm/build
|
|
cp /home/rocm/audio/dist/*.whl /home/rocm/build
|
|
EOF
|
|
|
|
RUN python${PYTHON_VERSION} -m pip install /home/rocm/build/*.whl
|
|
|
|
WORKDIR /home/rocm
|
|
CMD [ "bash", "-l" ]
|