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

[Anai-Guo]

Fixes #1946.

Problem

Linear4bit._save_to_state_dict writes the packed weight and the QuantState components:

• weight.absmax, weight.quant_map
• weight.nested_absmax, weight.nested_quant_map (when compress_statistics=True)
• weight.quant_state.bitsandbytes__nf4 / __fp4

But Linear4bit does not override _load_from_state_dict; it inherits nn.Linear._load_from_state_dict, which only consumes weight and bias. So:

• load_state_dict(strict=True) raises RuntimeError: Unexpected key(s) in state_dict for every QuantState entry.
• load_state_dict(strict=False) silently drops them. 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 — so forward passes return garbage. This is asymmetric vs Linear8bitLt, which already implements both halves (save at modules.py:1095, load at modules.py:1119).

Fix

Add Linear4bit._load_from_state_dict. After delegating to super(), walk unexpected_keys for entries under <prefix>weight., collect them into a qs_dict, rebuild via QuantState.from_dict(...), install on self.weight, and remove the consumed keys from unexpected_keys. Mirrors the existing Linear8bitLt pattern.

Reproducer (from the issue, abbreviated)

import bitsandbytes as bnb

src = bnb.nn.Linear4bit(64, 64, bias=False, quant_type='nf4', compute_dtype=torch.bfloat16, compress_statistics=True).to('cuda')
dst = bnb.nn.Linear4bit(64, 64, bias=False, quant_type='nf4', compute_dtype=torch.bfloat16, compress_statistics=True).to('cuda')

dst.load_state_dict(src.state_dict())
# before: RuntimeError: Unexpected key(s) ... weight.absmax, weight.quant_map, ...
# after:  dst.weight.quant_state matches src.weight.quant_state

Notes

• Behavior is unchanged for state_dicts that only contain weight (and bias): when no qs_prefix keys remain in unexpected_keys, the new code is a no-op early-return.
• No public API is changed; this only adds a private hook that PyTorch's load_state_dict machinery already expects modules to implement when they extend _save_to_state_dict.
• Embedding4bit has the same issue but a separate _save_to_state_dict shape; happy to extend the fix there in a follow-up if maintainers prefer one PR per class.

🤖 Generated with Claude Code

[matthewdouglas]

@SunMarc Could you help me understand if this is reasonable from the transformers/accelerate integration side?

[egeozkoc]

I spent some time testing this against the current HuggingFace stack to help answer the transformers/accelerate question, since that (plus the peft sensitivity) is what led to the original 2023 version being walked back in #864. Posting the results in case they help.

Environment: transformers==4.57.6, peft==0.19.1, accelerate==1.14.0, torch 2.12, CPU backend. Tiny model: hf-internal-testing/tiny-random-LlamaForCausalLM (14 Linear4bit layers).

What the bug actually costs

One thing worth flagging: the impact is broader than the top-level Linear4bit in the issue repro. It hits nested modules too, so basically any real model. With a Linear4bit nested under a prefix and no override:

load	result
strict=True	raises RuntimeError (5 QuantState keys unexpected), and forward output is garbage
strict=False	loads "successfully", silently drops the 5 keys, forward output is garbage

The strict=False case is the nasty one. The packed weight bytes get loaded but quant_state stays the stale one, so there's no error and the outputs are just wrong.

Integration paths (with this PR applied)

Scenario	_load_from_state_dict invoked?	Result
from_pretrained(load_in_4bit=True)	no (goes through the accelerate meta path)	works
save_pretrained then from_pretrained (serialized 4-bit)	no	logits match exactly
PEFT LoRA on 4-bit base, save_pretrained then PeftModel.from_pretrained	yes (per layer), 0 errors	logits match exactly
plain nn.Module.load_state_dict (nested), strict True/False	yes	logits match exactly
tests/test_linear4bit.py -k "serialization or copy or deepcopy"	n/a	164 passed

So on current versions the transformers and PEFT load paths route 4-bit weights through their own mechanisms, and they either skip this override entirely or hit it as a safe no-op. I couldn't get it to error or change behavior in any of those flows, and the peft#1095-style breakage doesn't seem to reproduce anymore. The override only does real work on the plain load_state_dict path, which is where the bug is.

One thing I couldn't rule out: a manual strict load onto a meta-device 4-bit model would run QuantState.from_dict(device=self.weight.device) against a meta device. I couldn't reach that through any normal transformers/PEFT flow (0 override calls there), so it's just untested rather than known good. Might be worth a guard or at least a note.

Suggested regression test

The existing test_linear_serialization restores via from_prequantized, so it never actually calls load_state_dict, which is probably why this went unnoticed. Here's a small test that covers the broken path. It fails on main and passes with this PR (I checked, 8/8 both ways):

class _Wrapper(torch.nn.Module):
    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):
    def build(seed):
        torch.manual_seed(seed)
        net = _Wrapper(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 is 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))

Happy to open a follow-up PR with this test (and a meta-device check) if that would help.