From c2abf7a8e34bc6ec12cb62bd5dc7ae95e89c70ec Mon Sep 17 00:00:00 2001 From: Afloat Date: Fri, 13 Nov 2026 04:19:03 +0800 Subject: [PATCH] Fix quantized Linear8bitLt checkpoint reload after forward --- bitsandbytes/nn/modules.py | 4 ++- tests/test_linear8bitlt.py | 57 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 1 deletion(-) diff --git a/bitsandbytes/nn/modules.py b/bitsandbytes/nn/modules.py index ebc0b0943..a0095782d 100644 --- a/bitsandbytes/nn/modules.py +++ b/bitsandbytes/nn/modules.py @@ -1141,6 +1141,8 @@ def _load_from_state_dict( input_name = key[len(prefix) :] if input_name == "SCB": weight_scb = getattr(self.weight, "SCB", None) + if weight_scb is None: + weight_scb = self.state.SCB if weight_scb is None: # buffers not yet initialized, can't access them directly without quantizing first raise RuntimeError( @@ -1152,7 +1154,7 @@ def _load_from_state_dict( weight_scb.copy_(input_param) if self.state.SCB is not None: - self.state.SCB = self.weight.SCB + self.state.SCB = weight_scb unexpected_keys.remove(key) diff --git a/tests/test_linear8bitlt.py b/tests/test_linear8bitlt.py index b078e82b7..ceba95ac5 100644 --- a/tests/test_linear8bitlt.py +++ b/tests/test_linear8bitlt.py @@ -249,6 +249,63 @@ def test_linear8bit_load_state_dict_raises_runtime_for_tied_weight(): linear.load_state_dict(state_dict, strict=False) +@pytest.mark.parametrize("device", get_available_devices()) +@pytest.mark.parametrize("threshold", [0.0, 6.0]) +@pytest.mark.parametrize("save_after_forward", TRUE_FALSE) +@pytest.mark.parametrize("destination_state", ["cold", "train", "eval"]) +def test_linear8bit_checkpoint_reload_after_forward(device, threshold, save_after_forward, destination_state): + torch.manual_seed(23) + destination = torch.nn.Sequential(Linear8bitLt(64, 32, has_fp16_weights=False, threshold=threshold)).to(device) + inputs = torch.randn(4, 64, device=device) + if destination_state != "cold": + destination.train(destination_state == "train") + with torch.no_grad(): + destination(inputs) + assert destination[0].weight.SCB is None + assert destination[0].state.SCB is not None + + for scale in (1.0, 3.0): + reference = torch.nn.Sequential(Linear8bitLt(64, 32, has_fp16_weights=False, threshold=threshold)) + with torch.no_grad(): + reference[0].weight.mul_(scale) + reference.to(device) + if save_after_forward: + with torch.no_grad(): + reference(inputs) + checkpoint = torch_load_from_buffer(torch_save_to_buffer(reference.state_dict())) + expected_checkpoint = copy.deepcopy(checkpoint) + destination.load_state_dict(checkpoint, strict=True) + torch.testing.assert_close(checkpoint, expected_checkpoint, rtol=0, atol=0) + + reference.eval() + destination.eval() + with torch.no_grad(): + torch.testing.assert_close(destination(inputs), reference(inputs), rtol=0, atol=0) + torch.testing.assert_close(destination[0].state.SCB, checkpoint["0.SCB"], rtol=0, atol=0) + assert destination[0].state.CB.data_ptr() == destination[0].weight.data_ptr() + torch.testing.assert_close(destination.state_dict(), checkpoint, rtol=0, atol=0) + + reference.train() + destination.train() + reference_optimizer = torch.optim.SGD([reference[0].bias], lr=0.1, momentum=0.9) + destination_optimizer = torch.optim.SGD([destination[0].bias], lr=0.1, momentum=0.9) + for _ in range(3): + reference_optimizer.zero_grad() + destination_optimizer.zero_grad() + reference_input = inputs.detach().clone().requires_grad_() + destination_input = inputs.detach().clone().requires_grad_() + expected = reference(reference_input) + actual = destination(destination_input) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + expected.square().mean().backward() + actual.square().mean().backward() + torch.testing.assert_close(destination_input.grad, reference_input.grad, rtol=0, atol=0) + torch.testing.assert_close(destination[0].bias.grad, reference[0].bias.grad, rtol=0, atol=0) + reference_optimizer.step() + destination_optimizer.step() + torch.testing.assert_close(destination[0].bias, reference[0].bias, rtol=0, atol=0) + + @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("threshold", [0.0, 6.0], ids=id_formatter("threshold")) @pytest.mark.parametrize("bias", TRUE_FALSE, ids=id_formatter("bias"))