From 07d79520efd7fc42ac9ce0a055f7ea6e583bc3a8 Mon Sep 17 00:00:00 2001
From: Ilyas Moutawwakil <57442720+IlyasMoutawwakil@users.noreply.github.com>
Date: Tue, 26 Mar 2024 14:43:12 +0100
Subject: [PATCH] Disable AMD memory benchmarks (#29871)

* remove py3nvml to skip amd memory benchmarks

* uninstall pynvml from docker images
---
 docker/transformers-pytorch-amd-gpu/Dockerfile           | 3 +++
 docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile | 5 ++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/docker/transformers-pytorch-amd-gpu/Dockerfile b/docker/transformers-pytorch-amd-gpu/Dockerfile
index 46ca1a531b..0b070c93a6 100644
--- a/docker/transformers-pytorch-amd-gpu/Dockerfile
+++ b/docker/transformers-pytorch-amd-gpu/Dockerfile
@@ -34,3 +34,6 @@ RUN python3 -m pip uninstall -y tensorflow flax
 # When installing in editable mode, `transformers` is not recognized as a package.
 # this line must be added in order for python to be aware of transformers.
 RUN cd transformers && python3 setup.py develop
+
+# Remove nvml as it is not compatible with ROCm
+RUN python3 -m pip uninstall py3nvml pynvml -y
diff --git a/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile b/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile
index 1fa384dfa2..fc6f912235 100644
--- a/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile
+++ b/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile
@@ -42,4 +42,7 @@ RUN python3 -m pip install --no-cache-dir ./transformers[accelerate,testing,sent
 # this line must be added in order for python to be aware of transformers.
 RUN cd transformers && python3 setup.py develop
 
-RUN python3 -c "from deepspeed.launcher.runner import main"
\ No newline at end of file
+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