fix(nn/Linear4bit): consume QuantState keys in _load_from_state_dict (#1946)

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

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))