Fix Params4bit.to() leaving CPU-packed weights packed

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:

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