Fix QuantState and dict conversions

[cyyever]

Fix the failure of QuantState.as_dict then followed by QuantState.from_dict.
A reproducing example is

import torch

from bitsandbytes.functional import (
    QuantState,
    dequantize_4bit,
    quantize_4bit,
)


a = torch.rand(2, 3)
quantized, quant_state = quantize_4bit(A=a, quant_type="nf4")
quant_dict = quant_state.as_dict()
quant_state = QuantState.from_dict(quant_dict, device=torch.device("cuda"))
quantized = dequantize_4bit(A=quantized, quant_state=quant_state)

Before this PR, it failed with (on bitsandbytes 0.47.0):

  File "/home/cyy/a.py", line 13, in <module>
    quant_state = QuantState.from_dict(quant_dict, device=torch.device("cuda"))
  File "/home/cyy/.local/lib/python3.13/site-packages/bitsandbytes/functional.py", line 469, in from_dict
    raise ValueError(
        f"There should be exactly one `quant_state` item with ending from {cls.valid_qs_type_keys}.\nDetected {qs_key}.",
    )
ValueError: There should be exactly one `quant_state` item with ending from ['bitsandbytes__fp4', 'bitsandbytes__nf4'].
Detected [].

[matthewdouglas]

Hi,
Can you provide more detail on the issue that this fixes? A simple reproducer with example output would be ideal. Thanks!

[cyyever]

@matthewdouglas Added

[github-actions]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[TimDettmers]

PR Review: #1729 — Fix QuantState and dict conversions

Bug fix: corrects two bugs in QuantState.as_dict() and one bug in QuantState.from_dict() that prevent the as_dict/from_dict round-trip from working when shape or quant_type are None (as produced by quantize_blockwise).

No blocking issues. The fix is correct, well-scoped, and includes regression tests.

Analysis

There are three distinct bugs being fixed:

1. as_dict() crashes when self.shape is None — quantize_blockwise() creates QuantState without setting shape (defaults to None). The old code unconditionally calls tuple(self.shape) on line 572, which raises TypeError. The fix adds a None guard: tuple(self.shape) if self.shape is not None else None.

2. as_dict(packed=True) crashes when self.quant_type is None — The old code unconditionally concatenates "bitsandbytes__" + self.quant_type on line 590, which raises TypeError when quant_type is None (again from quantize_blockwise). The fix conditionally appends the quant_type suffix only when it is not None.

3. from_dict() rejects valid unpacked dicts — The old validation logic at lines 523–528 used if/elif with not len(qs_key) and len(qs_key) != 1 checks. When given an unpacked dict (from as_dict(packed=False)) that has quant_type as a plain key but no quant_state.* packed key, the elif branch fires because len(qs_key) == 0 != 1, raising a misleading ValueError. The fix restructures the conditionals so that packed-format validation only runs when quant_type is absent from the dict, correctly recognizing unpacked dicts.

All three fixes are narrowly scoped and correct.

Serialization Compatibility

This is the critical concern for QuantState changes. After careful analysis:

• 4-bit checkpoints (the common case): quantize_4bit() always sets both shape and quant_type, so the as_dict(packed=True) path used by Linear4bit._save_to_state_dict() is unaffected. The packed key format quant_state.bitsandbytes__nf4 / quant_state.bitsandbytes__fp4 is unchanged.

• Old checkpoints loading with new code: Old checkpoints use the packed format with quant_type always set. from_dict() still handles these correctly — the new validation logic is equivalent to the old logic when a packed quant_state.* key is present.

• New checkpoints loading with old code: Not a concern for the packed format since the output is identical for the 4-bit case. For blockwise QuantState (which is never serialized to disk via _save_to_state_dict), the unpacked format now includes None values for shape and quant_type, but this dict is only used in-memory.

• Latent note: The as_dict(packed=True) path with quant_type=None produces a key "quant_state.bitsandbytes__" which from_dict would still reject (since "bitsandbytes__" is not in valid_qs_type_keys). This is not a practical issue because blockwise QuantState is never serialized via the packed path, but it is worth noting as a theoretical incomplete fix. Not blocking.

Test Coverage

The PR adds as_dict/from_dict round-trip coverage to four existing test functions:

• test_dynamic_blockwise_quantization (two loops: default code and custom code)
• test_fp8_quant
• test_4bit_quant

This covers both blockwise quantization (where shape/quant_type are None) and 4-bit quantization (where they are set). The tests use packed=False (the default), which is the path that was broken. The tests verify the round-trip by running dequantization after reconstruction, confirming the reconstructed QuantState is functionally correct.

The test does not cover the packed=True round-trip for blockwise QuantState, but as noted above, that code path is not used in practice.

Minor: Typo fix

The PR also fixes a documentation typo in the dequantize_4bit docstring: "The the absolute" -> "The absolute". This is fine.

• Security: Clear (no new imports, no network access, no filesystem writes, no obfuscation, no invisible Unicode characters)
• Downstream impact: None (4-bit checkpoint serialization format is unchanged; QuantState constructor signature unchanged; vLLM's QuantState.from_dict() usage is unaffected)
• Tests: Adequate (round-trip coverage for both blockwise and 4-bit paths)
• CI: Not triggered (fork PR — a maintainer needs to approve the workflow run before merge)
• Serialization: Compatible (no change to packed checkpoint format for 4-bit; blockwise QuantState is not persisted)
• Cross-PR conflicts: PR #1866 also modifies bitsandbytes/functional.py (adds __getattr__ to QuantState) but touches different lines. No direct conflict, though #1866's __getattr__ calls self.as_dict(packed=True) which would be affected by the as_dict change here. Since #1866 only exercises the 4-bit path (where quant_type is always set), there is no semantic conflict.
• Commit hygiene: 7 commits including 2 merge commits. Recommend squash merge.