Params4bit.to() skips the CPU AVX512 packing inverse that .cuda()/.xpu() apply, so model.to(device) silently corrupts weights

[caiotheodoro]

Params4bit.cuda() and .xpu() check quant_state.packing_format_for_cpu and undo the CPU AVX512 weight packing before moving the tensor:

def cuda(self, device=None, non_blocking=False):
    if getattr(self.quant_state, "packing_format_for_cpu", False):
        self.data, self.quant_state = _convert_weight_packed_for_cpu_inverse(self.data, self.quant_state)
    return self.to(device="cuda" if device is None else device, non_blocking=non_blocking)

Params4bit.to() (bitsandbytes/nn/modules.py, same class) doesn't do this check at all. It just calls torch._C._nn._parse_to and moves the data:

def to(self, *args, **kwargs):
    device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
    if device is not None and device.type != "meta" and not self.bnb_quantized:
        return self._quantize(device)
    else:
        if self.quant_state is not None:
            self.quant_state.to(device)
        new_param = Params4bit(super().to(device=device, dtype=dtype, non_blocking=non_blocking), ...)
        return new_param

nn.Module.to() moves parameters by calling Parameter.to(), not Parameter.cuda() (nn.Module.cuda() is the one that calls .cuda(), via self._apply(lambda t: t.cuda(device))). So the common model.to(device) / model.to("cuda") idiom — what HF transformers/accelerate and most user code actually calls — never hits the unpacking check. There's also no mps() method, so .to("mps") has no unpacking path at all, on either code path.

The packing (_convert_weight_packed_for_cpu, bitsandbytes/functional.py) is a real bit-layout transform for the AVX512 CPU kernel: it unpacks each nibble, transposes/regroups them into 32-row blocks, and repacks — not a view/reshape. Linear4bit.forward triggers it lazily the first time it runs on CPU with support_avx512bf16_for_cpu set. If that same module is later moved off CPU with model.to(device) (or .to("mps") at all, no matter what device it started on), the packed layout is never inverted, and any kernel dequantizing it after the move reads the wrong nibble order.

Repro, run on main at 8336490, torch 2.14.0, Python 3.13.12, macOS arm64 (packing_format_for_cpu is a plain state flag, so this reproduces without real AVX512 hardware):

import torch
from bitsandbytes.nn import Linear4bit
from bitsandbytes.functional import _convert_weight_packed_for_cpu, dequantize_4bit

torch.manual_seed(0)
lin = Linear4bit(64, 32, bias=False, quant_type="nf4", compute_dtype=torch.float32).to("cpu")
w0, qs0 = lin.weight.data.clone(), lin.weight.quant_state
ref = dequantize_4bit(w0, qs0)

# what Linear4bit.forward does on the first CPU AVX512 forward pass
packed_w, packed_qs = _convert_weight_packed_for_cpu(w0.clone(), qs0)
lin.weight.data, lin.weight.quant_state = packed_w, packed_qs

# the idiomatic call: nn.Module.to() -> Parameter.to(), not Parameter.cuda()
lin2 = lin.to("cpu")
out = dequantize_4bit(lin2.weight.data, lin2.weight.quant_state)
print("packing flag survives .to():", lin2.weight.quant_state.packing_format_for_cpu)
print("max abs diff vs pre-pack dequant:", (out - ref).abs().max().item())

packing flag survives .to(): True
max abs diff vs pre-pack dequant: 0.2499777376651764

0.25 on values that live in roughly the [-1, 1] NF4 range before scaling is not floating-point noise — the weight is silently wrong after the move, no error raised.

Fix: give Params4bit a to() that checks packing_format_for_cpu the same way cuda()/xpu() already do (or move the check into a shared helper all three call), and add an mps() override (or make the check device-agnostic in to() so it also covers .to("mps")).