AMD DeepSpeed image additional HIP dependencies (#36195)

* Add hipsolver and hipblastlt as dependencies

* Upgrade torch libs with rocm6.2.4 index
This commit is contained in:
ivarflakstad 2025-02-17 11:50:49 +01:00 committed by GitHub
parent 23d6095e8f
commit 936aeb70ab
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -2,10 +2,10 @@ FROM rocm/dev-ubuntu-22.04:6.2.4
LABEL maintainer="Hugging Face" LABEL maintainer="Hugging Face"
ARG DEBIAN_FRONTEND=noninteractive ARG DEBIAN_FRONTEND=noninteractive
ARG PYTORCH='2.5.1' ARG PYTORCH='2.6.0'
ARG TORCH_VISION='0.20.0' ARG TORCH_VISION='0.21.0'
ARG TORCH_AUDIO='2.5.0' ARG TORCH_AUDIO='2.6.0'
ARG ROCM='6.2' ARG ROCM='6.2.4'
RUN apt update && \ RUN apt update && \
apt install -y --no-install-recommends \ apt install -y --no-install-recommends \
@ -16,9 +16,11 @@ RUN apt update && \
python-is-python3 \ python-is-python3 \
rocrand-dev \ rocrand-dev \
rocthrust-dev \ rocthrust-dev \
rocblas-dev \
hipsolver-dev \
hipsparse-dev \ hipsparse-dev \
hipblas-dev \ hipblas-dev \
rocblas-dev && \ hipblaslt-dev && \
apt clean && \ apt clean && \
rm -rf /var/lib/apt/lists/* rm -rf /var/lib/apt/lists/*