Update CUDA/ROCm setup tests (#1899)

* Update CUDA/ROCm setup tests

* Fix pre-commit

* Add the assert that was removed by mistake

* Fix fragile override tests and restructure backend override logic

- Clear stray env vars in override tests to prevent false failures
- Branch on active backend first to handle both-overrides-set case
- Warn (instead of silently ignoring) the wrong override when both are set
- Reject overrides on unsupported backends (e.g. XPU)
- Add symmetric both-overrides-set tests for CUDA and ROCm

Author: sstamenk

bitsandbytes/cextension.py

@@ -30,29 +30,52 @@ def get_cuda_bnb_library_path(cuda_specs: CUDASpecs) -> Path:
     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")
+    cuda_override_value = os.environ.get("BNB_CUDA_VERSION")
     rocm_override_value = os.environ.get("BNB_ROCM_VERSION")
 
-    if rocm_override_value and torch.version.hip:
-        library_name = re.sub(r"rocm\d+", f"rocm{rocm_override_value}", library_name, count=1)
-        logger.warning(
-            f"WARNING: BNB_ROCM_VERSION={rocm_override_value} environment variable detected; loading {library_name}.\n"
-            "This can be used to load a bitsandbytes version built with a ROCm version that is different from the PyTorch ROCm version.\n"
-            "If this was unintended set the BNB_ROCM_VERSION variable to an empty string: export BNB_ROCM_VERSION=\n"
-        )
-    elif override_value:
-        library_name = re.sub(r"cuda\d+", f"cuda{override_value}", library_name, count=1)
-        if torch.version.hip:
+    if torch.version.hip:
+        if cuda_override_value:
+            if not rocm_override_value:
+                raise RuntimeError(
+                    f"BNB_CUDA_VERSION={cuda_override_value} detected but this is not a CUDA build!\n"
+                    "Use BNB_ROCM_VERSION instead: export BNB_ROCM_VERSION=<version>\n"
+                    "Clear the variable and retry: unset BNB_CUDA_VERSION\n"
+                )
+            logger.warning(
+                f"WARNING: BNB_CUDA_VERSION={cuda_override_value} is set but ignored on this ROCm build. "
+                "Clear the variable: unset BNB_CUDA_VERSION",
+            )
+        if rocm_override_value:
+            library_name = re.sub(r"rocm\d+", f"rocm{rocm_override_value}", library_name, count=1)
+            logger.warning(
+                f"WARNING: BNB_ROCM_VERSION={rocm_override_value} environment variable detected; loading {library_name}.\n"
+                "This can be used to load a bitsandbytes version built with a ROCm version that is different from the PyTorch ROCm version.\n"
+                "If this was unintended clear the variable and retry: unset BNB_ROCM_VERSION\n",
+            )
+    elif torch.version.cuda:
+        if rocm_override_value:
+            if not cuda_override_value:
+                raise RuntimeError(
+                    f"BNB_ROCM_VERSION={rocm_override_value} detected but this is not a ROCm build!\n"
+                    "Use BNB_CUDA_VERSION instead: export BNB_CUDA_VERSION=<version>\n"
+                    "Clear the variable and retry: unset BNB_ROCM_VERSION\n"
+                )
+            logger.warning(
+                f"WARNING: BNB_ROCM_VERSION={rocm_override_value} is set but ignored on this CUDA build. "
+                "Clear the variable: unset BNB_ROCM_VERSION",
+            )
+        if cuda_override_value:
+            library_name = re.sub(r"cuda\d+", f"cuda{cuda_override_value}", library_name, count=1)
+            logger.warning(
+                f"WARNING: BNB_CUDA_VERSION={cuda_override_value} environment variable detected; loading {library_name}.\n"
+                "This can be used to load a bitsandbytes version built with a CUDA version that is different from the PyTorch CUDA version.\n"
+                "If this was unintended clear the variable and retry: unset BNB_CUDA_VERSION\n",
+            )
+    else:
+        if rocm_override_value or cuda_override_value:
             raise RuntimeError(
-                f"BNB_CUDA_VERSION={override_value} detected for ROCm!! \n"
-                f"Use BNB_ROCM_VERSION instead: export BNB_ROCM_VERSION=<version>\n"
-                f"Clear the variable and retry: export BNB_CUDA_VERSION=\n"
+                "BNB_ROCM_VERSION / BNB_CUDA_VERSION overrides are not supported on this backend.",
             )
-        logger.warning(
-            f"WARNING: BNB_CUDA_VERSION={override_value} environment variable detected; loading {library_name}.\n"
-            "This can be used to load a bitsandbytes version built with a CUDA version that is different from the PyTorch CUDA version.\n"
-            "If this was unintended set the BNB_CUDA_VERSION variable to an empty string: export BNB_CUDA_VERSION=\n"
-        )
 
     return PACKAGE_DIR / library_name
 

tests/test_cuda_setup_evaluator.py

@@ -1,72 +1,98 @@
 import pytest
 
-from bitsandbytes.cextension import HIP_ENVIRONMENT, get_cuda_bnb_library_path
+from bitsandbytes.cextension import BNB_BACKEND, get_cuda_bnb_library_path
 from bitsandbytes.cuda_specs import CUDASpecs
 
 
 @pytest.fixture
 def cuda120_spec() -> CUDASpecs:
+    """Simulates torch+cuda12.0 and a representative Ampere-class capability."""
     return CUDASpecs(
         cuda_version_string="120",
         highest_compute_capability=(8, 6),
         cuda_version_tuple=(12, 0),
     )
 
 
-@pytest.mark.skipif(HIP_ENVIRONMENT, reason="this test is not supported on ROCm")
+@pytest.mark.skipif(BNB_BACKEND != "CUDA", reason="this test requires a CUDA backend")
 def test_get_cuda_bnb_library_path(monkeypatch, cuda120_spec):
+    """Without overrides, library path uses the detected CUDA 12.0 version."""
+    monkeypatch.delenv("BNB_ROCM_VERSION", raising=False)
     monkeypatch.delenv("BNB_CUDA_VERSION", raising=False)
     assert get_cuda_bnb_library_path(cuda120_spec).stem == "libbitsandbytes_cuda120"
 
 
-@pytest.mark.skipif(HIP_ENVIRONMENT, reason="this test is not supported on ROCm")
+@pytest.mark.skipif(BNB_BACKEND != "CUDA", reason="this test requires a CUDA backend")
 def test_get_cuda_bnb_library_path_override(monkeypatch, cuda120_spec, caplog):
+    """BNB_CUDA_VERSION=110 overrides path selection to the CUDA 11.0 binary."""
+    monkeypatch.delenv("BNB_ROCM_VERSION", raising=False)
     monkeypatch.setenv("BNB_CUDA_VERSION", "110")
     assert get_cuda_bnb_library_path(cuda120_spec).stem == "libbitsandbytes_cuda110"
     assert "BNB_CUDA_VERSION" in caplog.text  # did we get the warning?
 
 
-# Simulates torch+rocm7.0 (PyTorch bundled ROCm) on a system with ROCm 7.2
+@pytest.mark.skipif(BNB_BACKEND != "CUDA", reason="this test requires a CUDA backend")
+def test_get_cuda_bnb_library_path_rejects_rocm_override(monkeypatch, cuda120_spec):
+    """BNB_ROCM_VERSION alone should be rejected on CUDA with a helpful error."""
+    monkeypatch.delenv("BNB_CUDA_VERSION", raising=False)
+    monkeypatch.setenv("BNB_ROCM_VERSION", "72")
+    with pytest.raises(RuntimeError, match=r"BNB_ROCM_VERSION.*not a ROCm build"):
+        get_cuda_bnb_library_path(cuda120_spec)
+
+
+@pytest.mark.skipif(BNB_BACKEND != "CUDA", reason="this test requires a CUDA backend")
+def test_get_cuda_bnb_library_path_cuda_override_takes_priority(monkeypatch, cuda120_spec, caplog):
+    """When both overrides are set on CUDA, the CUDA override wins and the ROCm one is warned about."""
+    monkeypatch.setenv("BNB_CUDA_VERSION", "110")
+    monkeypatch.setenv("BNB_ROCM_VERSION", "72")
+    assert get_cuda_bnb_library_path(cuda120_spec).stem == "libbitsandbytes_cuda110"
+    assert "BNB_CUDA_VERSION" in caplog.text
+    assert "BNB_ROCM_VERSION" in caplog.text
+    assert "ignored on this CUDA build" in caplog.text
+
+
 @pytest.fixture
 def rocm70_spec() -> CUDASpecs:
+    """Simulates torch+rocm7.0 (bundled ROCm) when the system ROCm is newer."""
     return CUDASpecs(
-        cuda_version_string="70",  # from torch.version.hip == "7.0.x"
-        highest_compute_capability=(0, 0),  # unused for ROCm library path resolution
+        cuda_version_string="70",
+        highest_compute_capability=(0, 0),
         cuda_version_tuple=(7, 0),
     )
 
 
-@pytest.mark.skipif(not HIP_ENVIRONMENT, reason="this test is only supported on ROCm")
+@pytest.mark.skipif(BNB_BACKEND != "ROCm", reason="this test requires a ROCm backend")
 def test_get_rocm_bnb_library_path(monkeypatch, rocm70_spec):
     """Without override, library path uses PyTorch's ROCm 7.0 version."""
     monkeypatch.delenv("BNB_ROCM_VERSION", raising=False)
     monkeypatch.delenv("BNB_CUDA_VERSION", raising=False)
     assert get_cuda_bnb_library_path(rocm70_spec).stem == "libbitsandbytes_rocm70"
 
 
-@pytest.mark.skipif(not HIP_ENVIRONMENT, reason="this test is only supported on ROCm")
+@pytest.mark.skipif(BNB_BACKEND != "ROCm", reason="this test requires a ROCm backend")
 def test_get_rocm_bnb_library_path_override(monkeypatch, rocm70_spec, caplog):
     """BNB_ROCM_VERSION=72 overrides to load the ROCm 7.2 library instead of 7.0."""
-    monkeypatch.setenv("BNB_ROCM_VERSION", "72")
     monkeypatch.delenv("BNB_CUDA_VERSION", raising=False)
+    monkeypatch.setenv("BNB_ROCM_VERSION", "72")
     assert get_cuda_bnb_library_path(rocm70_spec).stem == "libbitsandbytes_rocm72"
     assert "BNB_ROCM_VERSION" in caplog.text
 
 
-@pytest.mark.skipif(not HIP_ENVIRONMENT, reason="this test is only supported on ROCm")
-def test_get_rocm_bnb_library_path_rejects_cuda_override(monkeypatch, rocm70_spec):
-    """BNB_CUDA_VERSION should be rejected on ROCm with a helpful error."""
-    monkeypatch.delenv("BNB_ROCM_VERSION", raising=False)
-    monkeypatch.setenv("BNB_CUDA_VERSION", "72")
-    with pytest.raises(RuntimeError, match=r"BNB_CUDA_VERSION.*detected for ROCm"):
-        get_cuda_bnb_library_path(rocm70_spec)
-
-
-@pytest.mark.skipif(not HIP_ENVIRONMENT, reason="this test is only supported on ROCm")
+@pytest.mark.skipif(BNB_BACKEND != "ROCm", reason="this test requires a ROCm backend")
 def test_get_rocm_bnb_library_path_rocm_override_takes_priority(monkeypatch, rocm70_spec, caplog):
-    """When both are set, BNB_ROCM_VERSION wins if HIP_ENVIRONMENT is True."""
+    """When both overrides are set on ROCm, the ROCm override wins and the CUDA one is warned about."""
     monkeypatch.setenv("BNB_ROCM_VERSION", "72")
-    monkeypatch.setenv("BNB_CUDA_VERSION", "72")
+    monkeypatch.setenv("BNB_CUDA_VERSION", "110")
     assert get_cuda_bnb_library_path(rocm70_spec).stem == "libbitsandbytes_rocm72"
     assert "BNB_ROCM_VERSION" in caplog.text
-    assert "BNB_CUDA_VERSION" not in caplog.text
+    assert "BNB_CUDA_VERSION" in caplog.text
+    assert "ignored on this ROCm build" in caplog.text
+
+
+@pytest.mark.skipif(BNB_BACKEND != "ROCm", reason="this test requires a ROCm backend")
+def test_get_rocm_bnb_library_path_rejects_cuda_override(monkeypatch, rocm70_spec):
+    """BNB_CUDA_VERSION alone should be rejected on ROCm with a helpful error."""
+    monkeypatch.delenv("BNB_ROCM_VERSION", raising=False)
+    monkeypatch.setenv("BNB_CUDA_VERSION", "110")
+    with pytest.raises(RuntimeError, match=r"BNB_CUDA_VERSION.*not a CUDA build"):
+        get_cuda_bnb_library_path(rocm70_spec)