Stub out for multi-platform

Author: matthewdouglas

bitsandbytes/__init__.py

@@ -4,18 +4,34 @@
 # LICENSE file in the root directory of this source tree.
 
 
+import torch
+
 from . import _ops, research, utils
 from .autograd._functions import (
     MatmulLtState,
     matmul,
     matmul_4bit,
 )
 from .backends.cpu import ops as cpu_ops
-from .backends.cuda import ops as cuda_ops  ## TODO: We would guard this for CUDA only
 from .backends.default import ops as default_ops
 from .nn import modules
 from .optim import adam
 
+# This is a signal for integrations with transformers/diffusers.
+# Eventually, we will remove this and check based on release version.
+features = {"multi-backend"}
+supported_torch_devices = {
+    "cuda",
+    "cpu",
+    # "mps",
+    # "xpu",
+    # "hpu",
+    # "npu",
+}
+
+if torch.cuda.is_available():
+    from .backends.cuda import ops as cuda_ops
+
 __pdoc__ = {
     "libbitsandbytes": False,
     "optim.optimizer.Optimizer8bit": False,

bitsandbytes/cextension.py

@@ -19,7 +19,9 @@ def get_cuda_bnb_library_path(cuda_specs: CUDASpecs) -> Path:
 
     The library is not guaranteed to exist at the returned path.
     """
-    library_name = f"libbitsandbytes_cuda{cuda_specs.cuda_version_string}{DYNAMIC_LIBRARY_SUFFIX}"
+
+    prefix = "rocm" if torch.version.hip else "cuda"
+    library_name = f"libbitsandbytes_{prefix}{cuda_specs.cuda_version_string}{DYNAMIC_LIBRARY_SUFFIX}"
 
     override_value = os.environ.get("BNB_CUDA_VERSION")
     if override_value:
@@ -76,7 +78,7 @@ def get_native_library() -> BNBNativeLibrary:
 
     logger.warning(
         "The installed version of bitsandbytes was compiled without GPU support. "
-        "8-bit optimizers, 8-bit multiplication, and GPU quantization are unavailable.",
+        "8-bit optimizers and GPU quantization are unavailable.",
     )
     return BNBNativeLibrary(dll)
 

bitsandbytes/cuda_specs.py

@@ -1,4 +1,5 @@
 import dataclasses
+from functools import lru_cache
 from typing import List, Optional, Tuple
 
 import torch
@@ -12,22 +13,26 @@ class CUDASpecs:
 
     @property
     def has_imma(self) -> bool:
-        return self.highest_compute_capability >= (7, 5)
+        return torch.version.hip or self.highest_compute_capability >= (7, 5)
 
 
 def get_compute_capabilities() -> List[Tuple[int, int]]:
     return sorted(torch.cuda.get_device_capability(torch.cuda.device(i)) for i in range(torch.cuda.device_count()))
 
 
+@lru_cache(None)
 def get_cuda_version_tuple() -> Tuple[int, int]:
-    # https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART____VERSION.html#group__CUDART____VERSION
-    major, minor = map(int, torch.version.cuda.split("."))
-    return major, minor
+    if torch.version.cuda:
+        return map(int, torch.version.cuda.split(".")[0:2])
+    elif torch.version.hip:
+        return map(int, torch.version.hip.split(".")[0:2])
+
+    return None
 
 
 def get_cuda_version_string() -> str:
     major, minor = get_cuda_version_tuple()
-    return f"{major}{minor}"
+    return f"{major * 10 + minor}"
 
 
 def get_cuda_specs() -> Optional[CUDASpecs]: