Skip to content

Commit

Permalink
Use rocm6.2 for AMD images (huggingface#35930)
Browse files Browse the repository at this point in the history
* Use rocm6.2 as rocm6.3 only has nightly pytorch wheels atm

* Use stable wheel index for torch libs
  • Loading branch information
ivarflakstad authored Jan 28, 2025
1 parent bf16a18 commit 96625d8
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 7 deletions.
6 changes: 2 additions & 4 deletions docker/transformers-pytorch-amd-gpu/Dockerfile
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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"

Expand Down
6 changes: 3 additions & 3 deletions docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile
Original file line number Diff line number Diff line change
@@ -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 \
Expand Down Expand Up @@ -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

0 comments on commit 96625d8

Please sign in to comment.