From 936aeb70abe14cfbc70fb4d4d0c7b1864a21cc54 Mon Sep 17 00:00:00 2001
From: ivarflakstad <69173633+ivarflakstad@users.noreply.github.com>
Date: Mon, 17 Feb 2025 11:50:49 +0100
Subject: [PATCH] AMD DeepSpeed image additional HIP dependencies (#36195)

* Add hipsolver and hipblastlt as dependencies

* Upgrade torch libs with rocm6.2.4 index
---
 .../Dockerfile                                       | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile b/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile
index a8edb8ff03e..f70b1549410 100644
--- a/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile
+++ b/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile
@@ -2,10 +2,10 @@ 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.2'
+ARG PYTORCH='2.6.0'
+ARG TORCH_VISION='0.21.0'
+ARG TORCH_AUDIO='2.6.0'
+ARG ROCM='6.2.4'
 
 RUN apt update && \
     apt install -y --no-install-recommends \
@@ -16,9 +16,11 @@ RUN apt update && \
     python-is-python3 \
     rocrand-dev \
     rocthrust-dev \
+    rocblas-dev \
+    hipsolver-dev \
     hipsparse-dev \
     hipblas-dev \
-    rocblas-dev && \
+    hipblaslt-dev && \
     apt clean && \
     rm -rf /var/lib/apt/lists/*