Fix _Linear8bitLt quantized guards and O(N^2) missing-keys hook in bitsandbytes plugin

src/lightning/fabric/plugins/precision/bitsandbytes.py

@@ -170,9 +170,7 @@ def _quantize_on_load_hook(quantize_fn: Callable[[torch.Tensor], None], state_di
 def _ignore_missing_weights_hook(module: torch.nn.Module, incompatible_keys: _IncompatibleKeys) -> None:
     # since we manually loaded the weight in the `_quantize_on_load_hook` hook, we need to avoid this missing key false
     # positive
-    for key in reversed(incompatible_keys.missing_keys):
-        if key.endswith("weight"):
-            incompatible_keys.missing_keys.remove(key)
+    incompatible_keys.missing_keys[:] = [key for key in incompatible_keys.missing_keys if not key.endswith("weight")]
 
 
 def _replace_param(param: torch.Tensor, data: torch.Tensor, quant_state: Any = None) -> torch.nn.Parameter:
@@ -270,7 +268,7 @@ def to_empty(self, *, device: Optional[_DEVICE], recurse: bool = True) -> Self:
             if self.weight.device.type == "meta":
                 # need custom logic if int8params is on meta device
                 raise NotImplementedError
-            if self.weight.dtype == torch.uint8:  # was quantized
+            if self.weight.dtype == torch.int8:  # was quantized
                 # need the original shape here
                 raise NotImplementedError
             device = torch.device(device)
@@ -290,8 +288,8 @@ def reset_parameters(self) -> None:
                 bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
                 init.uniform_(self.bias, -bound, bound)
 
-            linear_init_finished = isinstance(self.weight, bnb.nn.Params4bit)
-            if linear_init_finished and self.weight.dtype == torch.uint8:  # was quantized
+            linear_init_finished = isinstance(self.weight, bnb.nn.Int8Params)
+            if linear_init_finished and self.weight.dtype == torch.int8:  # was quantized
                 # need the original shape here
                 raise NotImplementedError
             weight = self.weight.data

tests/tests_fabric/plugins/precision/test_bitsandbytes.py

@@ -275,3 +275,51 @@ def forward(self, x):
     model.load_state_dict(state_dict)
     assert model.linear.weight.dtype == torch.uint8
     assert model.linear.weight.shape == (128, 1)
+
+
+def test_ignore_missing_weights_hook_and_linear8bit_guards(monkeypatch):
+    from torch.nn.modules.module import _IncompatibleKeys
+
+    from lightning.fabric.plugins.precision.bitsandbytes import _ignore_missing_weights_hook
+
+    missing_keys = ["layers.0.weight", "layers.0.bias", "layers.1.weight", "norm.bias"]
+    incompatible = _IncompatibleKeys(missing_keys, ["unexpected.key"])
+    _ignore_missing_weights_hook(torch.nn.Linear(2, 2), incompatible)
+    assert incompatible.missing_keys is missing_keys
+    assert incompatible.missing_keys == ["layers.0.bias", "norm.bias"]
+    assert incompatible.unexpected_keys == ["unexpected.key"]
+
+    module = lightning.fabric.plugins.precision.bitsandbytes
+    module._import_bitsandbytes.cache_clear()
+    monkeypatch.setattr(module, "_BITSANDBYTES_AVAILABLE", lambda: True)
+    bitsandbytes_mock = Mock()
+    monkeypatch.setitem(sys.modules, "bitsandbytes", bitsandbytes_mock)
+
+    class FakeInt8Params(torch.nn.Parameter):
+        pass
+
+    class FakeParams4bit(torch.nn.Parameter):
+        pass
+
+    class Module8bitMock(torch.nn.Linear):
+        def __init__(self, in_features, out_features, bias=True, *_, **__):
+            super().__init__(in_features, out_features, bias)
+            self.weight = FakeInt8Params(self.weight.data, requires_grad=False)
+
+    bitsandbytes_mock.nn.Linear8bitLt = Module8bitMock
+    bitsandbytes_mock.nn.Linear4bit = torch.nn.Linear
+    bitsandbytes_mock.nn.Int8Params = FakeInt8Params
+    bitsandbytes_mock.nn.Params4bit = FakeParams4bit
+
+    BitsandbytesPrecision("int8")
+    _Linear8bitLt = vars(module)["_Linear8bitLt"]
+    layer = _Linear8bitLt(4, 4, bias=False)
+    layer.weight = FakeInt8Params(torch.zeros((4, 4), dtype=torch.int8), requires_grad=False)
+
+    with pytest.raises(NotImplementedError):
+        layer.to_empty(device="cpu")
+
+    with pytest.raises(NotImplementedError):
+        layer.reset_parameters()
+
+    module._import_bitsandbytes.cache_clear()