diff --git a/bitsandbytes/nn/modules.py b/bitsandbytes/nn/modules.py index ebc0b0943..1f9d86949 100644 --- a/bitsandbytes/nn/modules.py +++ b/bitsandbytes/nn/modules.py @@ -606,6 +606,66 @@ def _save_to_state_dict(self, destination, prefix, keep_vars): for k, v in self.weight.quant_state.as_dict(packed=True).items(): destination[prefix + "weight." + k] = v if keep_vars else v.detach() + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + """ + Consume the QuantState components written by ``_save_to_state_dict``. + + Without this, the ``weight.absmax`` / ``weight.quant_map`` / ... keys + land in ``unexpected_keys`` and ``load_state_dict(strict=True)`` raises; + with ``strict=False`` they are silently dropped and the loaded layer + keeps a freshly-quantized ``quant_state`` that does not match the + packed bytes that were just loaded. + """ + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) + + qs_prefix = prefix + "weight." + qs_dict: dict[str, torch.Tensor] = {} + consumed_keys: list[str] = [] + for key in list(unexpected_keys): + if not key.startswith(qs_prefix): + continue + qs_dict[key[len(qs_prefix):]] = state_dict[key] + consumed_keys.append(key) + + if not qs_dict: + return + + try: + quant_state = QuantState.from_dict(qs_dict=qs_dict, device=self.weight.device) + except Exception as exc: + error_msgs.append( + f"Linear4bit: failed to reconstruct QuantState from state_dict " + f"with prefix '{prefix}': {exc}" + ) + return + + self.weight.quant_state = quant_state + self.weight.bnb_quantized = True + self.weight.blocksize = quant_state.blocksize + self.weight.compress_statistics = quant_state.nested + self.weight.quant_type = quant_state.quant_type + self.quant_state = quant_state + + for key in consumed_keys: + unexpected_keys.remove(key) + def forward(self, x: torch.Tensor): fix_4bit_weight_quant_state_from_module(self) quant_state = self.weight.quant_state diff --git a/tests/test_linear4bit.py b/tests/test_linear4bit.py index 12ed0eb27..fb36fd6db 100644 --- a/tests/test_linear4bit.py +++ b/tests/test_linear4bit.py @@ -615,3 +615,54 @@ def test_fsdp_state_dict_save_4bit(): f"stdout: {result.stdout}\n" f"stderr: {result.stderr}" ) + + +class _Linear4bitWrapper(torch.nn.Module): + """Nest Linear4bit under a prefix so load_state_dict exercises the real path.""" + + def __init__(self, quant_type, compress_statistics): + super().__init__() + self.fc = bnb.nn.Linear4bit( + 64, + 64, + bias=False, + compute_dtype=torch.float32, + compress_statistics=compress_statistics, + quant_type=quant_type, + ) + + +@pytest.mark.parametrize("device", get_available_devices()) +@pytest.mark.parametrize("quant_type", ["nf4", "fp4"]) +@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics")) +@pytest.mark.parametrize("strict", TRUE_FALSE, ids=id_formatter("strict")) +def test_linear4bit_load_from_state_dict(device, quant_type, compress_statistics, strict): + """Regression test for load_state_dict silently dropping QuantState keys. + + The existing serialization tests restore via ``from_prequantized`` and never + call ``load_state_dict``, so the ``_load_from_state_dict`` override was + uncovered. Without the fix, the packed ``weight`` bytes load but the stale + ``quant_state`` is kept: ``strict=True`` raises on unexpected keys and + ``strict=False`` silently produces wrong outputs. Fails on main, passes here. + """ + + def build(seed): + torch.manual_seed(seed) + net = _Linear4bitWrapper(quant_type, compress_statistics) + with torch.no_grad(): + net.fc.weight.data = torch.randn(64, 64) * 0.1 + return net.to(device) # triggers 4-bit quantization + + src = build(0) + dst = build(999) # different weights, so a no-op load would be caught + + result = dst.load_state_dict(src.state_dict(), strict=strict) + + assert result.missing_keys == [] + assert result.unexpected_keys == [] + assert dst.fc.weight.quant_state is not None + assert torch.equal(src.fc.weight.quant_state.absmax, dst.fc.weight.quant_state.absmax) + + x = torch.randn(8, 64, device=device) + with torch.no_grad(): + assert torch.equal(src.fc(x), dst.fc(x))