From 8dcedafdd61ed6a2de6f2002ea1f4bdd130d3b09 Mon Sep 17 00:00:00 2001 From: ailuntz Date: Tue, 10 Mar 2026 16:07:29 +0800 Subject: [PATCH 1/2] Allow ROCm arch override --- bitsandbytes/cuda_specs.py | 9 +++++++++ docs/source/installation.mdx | 3 +++ 2 files changed, 12 insertions(+) diff --git a/bitsandbytes/cuda_specs.py b/bitsandbytes/cuda_specs.py index f9fdd295d..86741943f 100644 --- a/bitsandbytes/cuda_specs.py +++ b/bitsandbytes/cuda_specs.py @@ -1,6 +1,7 @@ import dataclasses from functools import lru_cache import logging +import os import platform import re import subprocess @@ -83,6 +84,14 @@ def get_rocm_gpu_arch() -> str: """Get ROCm GPU architecture.""" logger = logging.getLogger(__name__) try: + override = os.environ.get("BNB_ROCM_GPU_ARCH") or os.environ.get("BNB_ROCM_ARCH") + if override: + override = override.strip() + if override: + if not override.startswith("gfx"): + override = f"gfx{override}" + logger.info("Using ROCm GPU arch override: %s", override) + return override if torch.version.hip: # On Windows, use hipinfo.exe; on Linux, use rocminfo if platform.system() == "Windows": diff --git a/docs/source/installation.mdx b/docs/source/installation.mdx index c97996b75..47e20ec7b 100644 --- a/docs/source/installation.mdx +++ b/docs/source/installation.mdx @@ -189,6 +189,9 @@ pip install -e . * All features are supported for both consumer RDNA devices and Data Center CDNA products. * A compatible PyTorch version with AMD ROCm support is required. It is recommended to use the latest stable release. See [PyTorch on ROCm](https://rocm.docs.amd.com/projects/install-on-linux/en/latest/install/3rd-party/pytorch-install.html) for guidance. +> [!TIP] +> If `rocminfo` is unavailable in your runtime environment, set `BNB_ROCM_GPU_ARCH` (for example `BNB_ROCM_GPU_ARCH=gfx90a`) to skip automatic detection. + ### Installation from PyPI[[rocm-pip]] This is the most straightforward and recommended installation option. From a19580173fcf9b483cd18de99d8833bb62fdc311 Mon Sep 17 00:00:00 2001 From: ailuntz Date: Sat, 14 Mar 2026 15:08:08 +0800 Subject: [PATCH 2/2] rocm arch override fallback --- bitsandbytes/cuda_specs.py | 24 ++++++++++++++++-------- docs/source/installation.mdx | 2 +- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/bitsandbytes/cuda_specs.py b/bitsandbytes/cuda_specs.py index 86741943f..3db42d536 100644 --- a/bitsandbytes/cuda_specs.py +++ b/bitsandbytes/cuda_specs.py @@ -84,14 +84,6 @@ def get_rocm_gpu_arch() -> str: """Get ROCm GPU architecture.""" logger = logging.getLogger(__name__) try: - override = os.environ.get("BNB_ROCM_GPU_ARCH") or os.environ.get("BNB_ROCM_ARCH") - if override: - override = override.strip() - if override: - if not override.startswith("gfx"): - override = f"gfx{override}" - logger.info("Using ROCm GPU arch override: %s", override) - return override if torch.version.hip: # On Windows, use hipinfo.exe; on Linux, use rocminfo if platform.system() == "Windows": @@ -106,6 +98,14 @@ def get_rocm_gpu_arch() -> str: if match: return "gfx" + match.group(1) else: + override = os.environ.get("BNB_ROCM_GPU_ARCH") + if override: + override = override.strip() + if override: + if not override.startswith("gfx"): + override = f"gfx{override}" + logger.info("Using ROCm GPU arch override: %s", override) + return override return "unknown" else: return "unknown" @@ -117,6 +117,14 @@ def get_rocm_gpu_arch() -> str: ROCm GPU architecture detection failed despite ROCm being available. """, ) + override = os.environ.get("BNB_ROCM_GPU_ARCH") + if override: + override = override.strip() + if override: + if not override.startswith("gfx"): + override = f"gfx{override}" + logger.info("Using ROCm GPU arch override: %s", override) + return override return "unknown" diff --git a/docs/source/installation.mdx b/docs/source/installation.mdx index 47e20ec7b..61c143b37 100644 --- a/docs/source/installation.mdx +++ b/docs/source/installation.mdx @@ -190,7 +190,7 @@ pip install -e . * A compatible PyTorch version with AMD ROCm support is required. It is recommended to use the latest stable release. See [PyTorch on ROCm](https://rocm.docs.amd.com/projects/install-on-linux/en/latest/install/3rd-party/pytorch-install.html) for guidance. > [!TIP] -> If `rocminfo` is unavailable in your runtime environment, set `BNB_ROCM_GPU_ARCH` (for example `BNB_ROCM_GPU_ARCH=gfx90a`) to skip automatic detection. +> If `rocminfo` is unavailable or fails in your runtime environment, set `BNB_ROCM_GPU_ARCH` (for example `BNB_ROCM_GPU_ARCH=gfx90a`) as a fallback. ### Installation from PyPI[[rocm-pip]]