Use rocm6.2 for AMD images (#35930)

* Use rocm6.2 as rocm6.3 only has nightly pytorch wheels atm

* Use stable wheel index for torch libs
This commit is contained in:
ivarflakstad 2025-01-28 11:10:28 +01:00 committed by GitHub
parent bf16a182ba
commit 96625d85fd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 5 additions and 7 deletions

View File

@ -1,4 +1,4 @@
FROM rocm/dev-ubuntu-22.04:6.3
FROM rocm/dev-ubuntu-22.04:6.2.4
LABEL maintainer="Hugging Face"
ARG DEBIAN_FRONTEND=noninteractive
@ -8,11 +8,9 @@ RUN apt update && \
apt clean && \
rm -rf /var/lib/apt/lists/*
RUN export PATH="${PATH:+${PATH}:}~/opt/rocm/bin"
RUN python3 -m pip install --no-cache-dir --upgrade pip numpy
RUN python3 -m pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.3/
RUN python3 -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.2
RUN python3 -m pip install --no-cache-dir --upgrade importlib-metadata setuptools ninja git+https://github.com/facebookresearch/detectron2.git pytesseract "itsdangerous<2.1.0"

View File

@ -1,11 +1,11 @@
FROM rocm/dev-ubuntu-22.04:6.3
FROM rocm/dev-ubuntu-22.04:6.2.4
LABEL maintainer="Hugging Face"
ARG DEBIAN_FRONTEND=noninteractive
ARG PYTORCH='2.5.1'
ARG TORCH_VISION='0.20.0'
ARG TORCH_AUDIO='2.5.0'
ARG ROCM='6.3'
ARG ROCM='6.2'
RUN apt update && \
apt install -y --no-install-recommends \
@ -45,4 +45,4 @@ RUN cd transformers && python3 setup.py develop
RUN python3 -c "from deepspeed.launcher.runner import main"
# Remove nvml as it is not compatible with ROCm
RUN python3 -m pip uninstall py3nvml pynvml -y
RUN python3 -m pip uninstall py3nvml pynvml nvidia-ml-py apex -y