fix broken cuda and rocm images

docker/cuda/Dockerfile

@@ -28,14 +28,18 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 
 # Install PyTorch
 ARG TORCH_CUDA=cu124
-ARG TORCH_VERSION=stable
+ARG TORCH_RELEASE_TYPE=stable
+ARG TORCH_VERSION
 
-RUN if [ "${TORCH_VERSION}" = "stable" ]; then \
-    pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
-elif [ "${TORCH_VERSION}" = "nighly" ]; then \
+RUN if [ -n "${TORCH_VERSION}" ]; then \
+    # Install specific torch version if TORCH_VERSION is provided
+    pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
+elif [ "${TORCH_RELEASE_TYPE}" = "stable" ]; then \
+    pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/${TORCH_CUDA} ; \
+elif [ "${TORCH_RELEASE_TYPE}" = "nightly" ]; then \
     pip install --no-cache-dir --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
 else \
-    pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
+    echo "Error: Invalid TORCH_RELEASE_TYPE. Must be 'stable', 'nightly', or specify a TORCH_VERSION." && exit 1 ; \
 fi
 
 # Install quantization libraries from source

docker/rocm/Dockerfile

@@ -31,16 +31,19 @@ RUN apt-get update && apt-get upgrade -y && apt-get install -y --no-install-reco
 
 # Install PyTorch
 ARG TORCH_ROCM=rocm5.7
-ARG TORCH_VERSION=stable
+ARG TORCH_RELEASE_TYPE=stable
+ARG TORCH_VERSION
 
-RUN if [ "${TORCH_VERSION}" = "stable" ]; then \
-    pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_ROCM} ; \
-elif [ "${TORCH_VERSION}" = "nightly" ]; then \
-    pip install --no-cache-dir --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/${TORCH_ROCM} ; \
-else \
+RUN if [ -n "${TORCH_VERSION}" ]; then \
+    # Install specific torch version if TORCH_VERSION is provided
     pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_ROCM} ; \
+elif [ "${TORCH_RELEASE_TYPE}" = "stable" ]; then \
+    pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/${TORCH_ROCM} ; \
+elif [ "${TORCH_RELEASE_TYPE}" = "nightly" ]; then \
+    pip install --no-cache-dir --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_ROCM} ; \
+else \
+    echo "Error: Invalid TORCH_RELEASE_TYPE. Must be 'stable', 'nightly', or specify a TORCH_VERSION." && exit 1 ; \
 fi
-
 # Install quantization libraries from source
 ENV ROCM_VERSION=5.7
 ENV PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx1030;gfx1100"

optimum_benchmark/hub_utils.py

@@ -9,7 +9,7 @@
 import pandas as pd
 from flatten_dict import flatten, unflatten
 from huggingface_hub import create_repo, hf_hub_download, upload_file
-from huggingface_hub.utils._errors import HfHubHTTPError
+from huggingface_hub.utils import HfHubHTTPError
 from typing_extensions import Self
 
 LOGGER = getLogger("hub_utils")

scripts/install_quantization_libs.py

@@ -1,29 +1,11 @@
 import argparse
 import os
-import re
 import subprocess
 import sys
 
 EXTERNAL_REPOS_DIR = "external_repos"
 
 
-def process_setup_file(setup_file_path):
-    with open(setup_file_path, "r") as file:
-        setup_content = file.read()
-
-    # Use a regular expression to remove any line containing "torch=="
-    setup_content = re.sub(r'"torch==[^\"]+",', "", setup_content)
-
-    # Set IS_CPU_ONLY to False
-    setup_content = setup_content.replace(
-        "IS_CPU_ONLY = not torch.backends.mps.is_available() and not torch.cuda.is_available()", "IS_CPU_ONLY = False"
-    )
-
-    # Write the modified content back to setup.py
-    with open(setup_file_path, "w") as file:
-        file.write(setup_content)
-
-
 def clone_or_pull_repo(repo_url, repo_location_path):
     """Clone the repo if it doesn't exist; otherwise, pull the latest changes."""
     if os.path.exists(repo_location_path):
@@ -36,36 +18,14 @@ def clone_or_pull_repo(repo_url, repo_location_path):
 
 
 def install_autoawq_from_source():
-    """Install the AutoAWQ and AutoAWQ_kernels packages from GitHub."""
-    print("Installing AutoAWQ and AutoAWQ_kernels packages.")
-
-    autoawq_repo_name = "AutoAWQ"
-    autoawq_kernels_repo_name = "AutoAWQ_kernels"
-
-    autoawq_repo_path = os.path.join(EXTERNAL_REPOS_DIR, autoawq_repo_name)
-    kernels_repo_path = os.path.join(EXTERNAL_REPOS_DIR, autoawq_kernels_repo_name)
-
-    clone_or_pull_repo(f"https://github.com/casper-hansen/{autoawq_kernels_repo_name}", kernels_repo_path)
-    kernels_setup_file_path = os.path.join(kernels_repo_path, "setup.py")
-    process_setup_file(kernels_setup_file_path)
-    subprocess.run(
-        f"cd {kernels_repo_path} && {sys.executable} -m pip install .",
-        shell=True,
-        check=True,
-        env=os.environ,
-    )
-
-    clone_or_pull_repo(f"https://github.com/casper-hansen/{autoawq_repo_name}", autoawq_repo_path)
-    autoawq_setup_file_path = os.path.join(autoawq_repo_path, "setup.py")
-    process_setup_file(autoawq_setup_file_path)
+    print("Installing AutoAWQ package.")
     subprocess.run(
-        f"cd {autoawq_repo_path} && {sys.executable} -m pip install .",
+        "INSTALL_KERNELS=1 pip install git+https://github.com/casper-hansen/AutoAWQ.git",
         shell=True,
         check=True,
         env=os.environ,
     )
-
-    print("AutoAWQ and AutoAWQ_kernels packages installed.")
+    print("AutoAWQ package installed.")
 
 
 def install_autogptq_from_source():
@@ -75,8 +35,6 @@ def install_autogptq_from_source():
 
     clone_or_pull_repo("https://github.com/PanQiWei/AutoGPTQ.git", autogptq_repo_path)
     subprocess.run("pip install numpy gekko pandas", shell=True, check=True, env=os.environ)
-    autogptq_setup_file_path = os.path.join(autogptq_repo_path, "setup.py")
-    process_setup_file(autogptq_setup_file_path)
     subprocess.run(
         f"cd {autogptq_repo_path} && {sys.executable} -m pip install .",
         shell=True,