From 40682fe3bda5cf2765842fa41c875311acdee48e Mon Sep 17 00:00:00 2001 From: 22elix3r Date: Thu, 10 Sep 2026 21:15:08 +0530 Subject: [PATCH] Fix Params4bit.to() leaving CPU-packed weights packed cuda() and xpu() undo AVX512 CPU packing before the move. nn.Module.to() calls Parameter.to(), so model.to(device) skipped that step and kept the packed nibble layout. Fixes #2078 --- bitsandbytes/nn/modules.py | 12 ++++++++++++ tests/test_linear4bit.py | 25 +++++++++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/bitsandbytes/nn/modules.py b/bitsandbytes/nn/modules.py index ebc0b0943..00fb5b64d 100644 --- a/bitsandbytes/nn/modules.py +++ b/bitsandbytes/nn/modules.py @@ -424,6 +424,18 @@ def to(self: T, tensor: Tensor, non_blocking: bool = ...) -> T: ... def to(self, *args, **kwargs): device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs) + # cuda()/xpu() already undo CPU AVX512 packing before calling to(). + # nn.Module.to() goes through Parameter.to(), so that check has to live + # here as well or model.to("cuda") keeps the packed nibble layout. + dest_type = device.type if device is not None else self.device.type + if ( + getattr(self.quant_state, "packing_format_for_cpu", False) + and dest_type not in ("cpu", "meta") + ): + self.data, self.quant_state = _convert_weight_packed_for_cpu_inverse( + self.data, self.quant_state + ) + if device is not None and device.type != "meta" and not self.bnb_quantized: return self._quantize(device) else: diff --git a/tests/test_linear4bit.py b/tests/test_linear4bit.py index 79ede45b2..4560a5a29 100644 --- a/tests/test_linear4bit.py +++ b/tests/test_linear4bit.py @@ -280,6 +280,31 @@ def test_quant_storage_shard_roundtrip(device, quant_type, quant_storage): torch.testing.assert_close(out, ref) +@pytest.mark.parametrize("device", get_available_devices()) +@pytest.mark.parametrize("quant_type", ["nf4", "fp4"]) +def test_params4bit_to_unpacks_cpu_packing(device, quant_type): + """Params4bit.to() must undo CPU AVX512 packing when leaving CPU (#2078).""" + if device == "cpu": + pytest.skip("Unpacking is only required when moving off CPU.") + if device == "hpu" and not is_supported_on_hpu(quant_type, torch.float32, torch.uint8): + pytest.skip("This configuration is not supported on HPU.") + + torch.manual_seed(0) + tensor = torch.randn(64, 32, dtype=torch.float32) + param = bnb.nn.Params4bit(data=tensor, quant_type=quant_type, requires_grad=False) + param = param._quantize("cpu") + ref = bnb.functional.dequantize_4bit(param.data.clone(), param.quant_state) + + packed_w, packed_qs = bnb.functional._convert_weight_packed_for_cpu(param.data.clone(), param.quant_state) + param.data, param.quant_state = packed_w, packed_qs + assert param.quant_state.packing_format_for_cpu + + moved = param.to(device) + assert not getattr(moved.quant_state, "packing_format_for_cpu", False) + out = bnb.functional.dequantize_4bit(moved.data, moved.quant_state) + torch.testing.assert_close(out.cpu().float(), ref.float(), atol=1e-5, rtol=1e-4) + + @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("quant_type", ["nf4", "fp4"]) @pytest.mark.parametrize("blocksize", [32, 64, 128])