Conversation
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 bitsandbytes-foundation#2078
There was a problem hiding this comment.
🟢 Approval recommended
The change is narrowly scoped, matches existing .cuda()/.xpu() behavior, and includes a targeted regression test covering the reported corruption scenario.
Pull request overview
This PR fixes a correctness bug in Params4bit.to() when moving a CPU-packed (AVX512) 4-bit weight to a non-CPU device via the common nn.Module.to(...) pathway. It ensures the CPU-specific packed nibble layout is inverted before transferring to accelerators, matching the behavior already present in .cuda()/.xpu() and preventing silent weight corruption.
Changes:
- Add a CPU-packing inverse step inside
Params4bit.to()when the destination device is not CPU/meta. - Add a regression test that packs a quantized CPU weight, moves it off CPU via
.to(device), and checks dequantized values match the pre-pack reference.
File summaries
| File | Description |
|---|---|
bitsandbytes/nn/modules.py |
Ensures Params4bit.to() undoes AVX512 CPU packing when moving to non-CPU devices, covering model.to(device) flows. |
tests/test_linear4bit.py |
Adds regression coverage for leaving CPU with packed 4-bit weights via .to(...). |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 0
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
@matthewdouglas @TimDettmers Params4bit.to() was skipping the CPU packing inverse that cuda()/xpu() already do, so model.to(device) could keep packed weights. Small patch plus a test. Review welcome. |
Params4bit.cuda()and.xpu()undo AVX512 CPU packing before the move.nn.Module.to()callsParameter.to(), somodel.to(device)skipped that step and kept the packed nibble layout.This moves the same unpack into
Params4bit.to()when the destination is not CPU.Fixes #2078