From 50b2b483f1dbf9412c1e38b2f92637c25d663143 Mon Sep 17 00:00:00 2001 From: Tai An Date: Thu, 28 May 2026 09:27:39 -0700 Subject: [PATCH 1/2] fix(nn/Linear4bit): consume QuantState keys in _load_from_state_dict (#1946) Linear4bit overrides _save_to_state_dict to write weight.absmax / weight.quant_map / weight.nested_* / weight.quant_state.bitsandbytes__* alongside the packed weight, but inherits nn.Linear._load_from_state_dict which only consumes weight and bias. Result: - strict=True load raises Unexpected key(s) in state_dict for every QuantState component. - strict=False silently drops them and the destination layer keeps the freshly-quantized quant_state from the prior .to('cuda') call, which does not match the packed bytes that were just loaded. This mirrors what Linear8bitLt already does for SCB (_load_from_state_dict at modules.py:1119): walk unexpected_keys for entries under 'weight.', collect them into a qs_dict, reconstruct via QuantState.from_dict, install on self.weight, and remove the consumed keys from unexpected_keys. Fixes #1946 --- bitsandbytes/nn/modules.py | 60 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) 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 From 5d620d91ac42690a0453abbf79239d89395248cd Mon Sep 17 00:00:00 2001 From: Tai An Date: Fri, 10 Jul 2026 21:10:00 -0700 Subject: [PATCH 2/2] test(linear4bit): cover load_state_dict QuantState path (#1907) Add a regression test for the nested Linear4bit load_state_dict path that the fix repairs. The existing serialization tests restore via from_prequantized and never call load_state_dict, so the override was uncovered. This test fails on main (dropped QuantState keys, garbage forward output) and passes with the fix, for both strict values. Test authored by @egeozkoc during PR review. --- tests/test_linear4bit.py | 51 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) 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))