Error message to inform bitsandbytes is only supported on CUDA (#19360)

Author: awaelchli

src/lightning/fabric/connector.py

@@ -24,6 +24,7 @@
 from lightning.fabric.accelerators.mps import MPSAccelerator
 from lightning.fabric.accelerators.xla import XLAAccelerator
 from lightning.fabric.plugins import (
+    BitsandbytesPrecision,
     CheckpointIO,
     DeepSpeedPrecision,
     HalfPrecision,
@@ -448,6 +449,10 @@ def _init_strategy(self) -> None:
 
     def _check_and_init_precision(self) -> Precision:
         if isinstance(self._precision_instance, Precision):
+            if isinstance(self._precision_instance, BitsandbytesPrecision) and not isinstance(
+                self.accelerator, CUDAAccelerator
+            ):
+                raise RuntimeError("Bitsandbytes is only supported on CUDA GPUs.")
             return self._precision_instance
         if isinstance(self.strategy, (SingleDeviceXLAStrategy, XLAStrategy, XLAFSDPStrategy)):
             return XLAPrecision(self._precision_input)  # type: ignore

src/lightning/pytorch/trainer/connectors/accelerator_connector.py

@@ -37,6 +37,7 @@
 from lightning.pytorch.accelerators.xla import XLAAccelerator
 from lightning.pytorch.plugins import (
     _PLUGIN_INPUT,
+    BitsandbytesPrecision,
     CheckpointIO,
     DeepSpeedPrecision,
     DoublePrecision,
@@ -565,6 +566,10 @@ def _check_and_init_precision(self) -> Precision:
 
     def _validate_precision_choice(self) -> None:
         """Validate the combination of choices for precision, AMP type, and accelerator."""
+        if isinstance(self._precision_plugin_flag, BitsandbytesPrecision) and not isinstance(
+            self.accelerator, CUDAAccelerator
+        ):
+            raise RuntimeError("Bitsandbytes is only supported on CUDA GPUs.")
         if _habana_available_and_importable():
             from lightning_habana import HPUAccelerator
 

tests/tests_fabric/test_connector.py

@@ -13,6 +13,7 @@
 # limitations under the License
 import inspect
 import os
+import sys
 from typing import Any, Dict
 from unittest import mock
 from unittest.mock import Mock
@@ -29,6 +30,7 @@
 from lightning.fabric.accelerators.mps import MPSAccelerator
 from lightning.fabric.connector import _Connector
 from lightning.fabric.plugins import (
+    BitsandbytesPrecision,
     DeepSpeedPrecision,
     DoublePrecision,
     FSDPPrecision,
@@ -864,6 +866,13 @@ def test_precision_selection_amp_ddp(strategy, devices, is_custom_plugin, plugin
     assert isinstance(connector.precision, plugin_cls)
 
 
+def test_bitsandbytes_precision_cuda_required(monkeypatch):
+    monkeypatch.setattr(lightning.fabric.plugins.precision.bitsandbytes, "_BITSANDBYTES_AVAILABLE", True)
+    monkeypatch.setitem(sys.modules, "bitsandbytes", Mock())
+    with pytest.raises(RuntimeError, match="Bitsandbytes is only supported on CUDA GPUs"):
+        _Connector(accelerator="cpu", plugins=BitsandbytesPrecision(mode="int8"))
+
+
 @pytest.mark.parametrize(("strategy", "strategy_cls"), [("DDP", DDPStrategy), ("Ddp", DDPStrategy)])
 @mock.patch("lightning.fabric.accelerators.mps.MPSAccelerator.is_available", return_value=False)
 def test_strategy_str_passed_being_case_insensitive(_, strategy, strategy_cls):

tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py

@@ -18,6 +18,7 @@
 from unittest import mock
 from unittest.mock import Mock
 
+import lightning.fabric
 import lightning.pytorch
 import pytest
 import torch
@@ -35,6 +36,7 @@
 from lightning.pytorch.plugins.io import TorchCheckpointIO
 from lightning.pytorch.plugins.layer_sync import LayerSync, TorchSyncBatchNorm
 from lightning.pytorch.plugins.precision import (
+    BitsandbytesPrecision,
     DeepSpeedPrecision,
     DoublePrecision,
     FSDPPrecision,
@@ -1115,3 +1117,10 @@ def test_connector_num_nodes_input_validation():
 def test_precision_selection(precision_str, strategy_str, expected_precision_cls):
     connector = _AcceleratorConnector(precision=precision_str, strategy=strategy_str)
     assert isinstance(connector.precision_plugin, expected_precision_cls)
+
+
+def test_bitsandbytes_precision_cuda_required(monkeypatch):
+    monkeypatch.setattr(lightning.fabric.plugins.precision.bitsandbytes, "_BITSANDBYTES_AVAILABLE", True)
+    monkeypatch.setitem(sys.modules, "bitsandbytes", Mock())
+    with pytest.raises(RuntimeError, match="Bitsandbytes is only supported on CUDA GPUs"):
+        _AcceleratorConnector(accelerator="cpu", plugins=BitsandbytesPrecision(mode="int8"))