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

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:

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)

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.

Thanks @egeozkoc, this is really helpful — especially confirming that the transformers/PEFT/accelerate load paths route 4-bit weights through their own mechanisms and hit this override either not at all or as a safe no-op, so the peft#1095-style breakage from #864 doesn't reproduce on current versions. That was exactly the open question.

I've added your regression test to tests/test_linear4bit.py (credited to you in the commit message). It nests Linear4bit under a prefix so load_state_dict actually runs the override, uses different source/destination weights so a no-op load is caught, and checks both strict values — fails on main, passes here.

Good call on the meta-device strict load; I left it out for now since it isn't reachable through any normal transformers/PEFT flow (0 override calls there), so it'd be testing an untested-but-unused path. Happy to add a guard + note in a follow-up if maintainers prefer.