Fix Params4bit/QuantState attribute access for FSDP state_dict traversal (#1405)

PyTorch's FSDP state_dict machinery (_get_fqns) resolves dotted FQN paths
via getattr. For 4-bit quantized models, state_dict keys like
"weight.absmax" and "weight.quant_state.bitsandbytes__nf4" require attribute
access on Params4bit and QuantState objects that previously didn't exist.

Add __getattr__ to Params4bit that proxies known QuantState attributes
(including the quant_map→code alias used by as_dict serialization), and
add __getattr__ to QuantState that handles the packed "bitsandbytes__*"
keys. This allows FSDP's get_model_state_dict with cpu_offload=True to
traverse the full FQN namespace without AttributeError.

Verified with single-GPU FSDP integration test: without fix, fails with
"'Params4bit' object has no attribute 'absmax'"; with fix, successfully
produces all 7 state_dict keys.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: TimDettmers

bitsandbytes/functional.py

@@ -487,6 +487,18 @@ def __init__(
         self.state2 = state2
         self.nested = state2 is not None
 
+    def __getattr__(self, name):
+        # Support attribute access for packed state_dict keys like "bitsandbytes__nf4".
+        # PyTorch's FSDP state_dict traversal (_get_fqns) resolves dotted FQN paths via
+        # getattr. The packed key "quant_state.bitsandbytes__nf4" causes it to call
+        # getattr(quant_state_obj, "bitsandbytes__nf4"), which we handle here.
+        if name.startswith("bitsandbytes__"):
+            qs_dict = self.as_dict(packed=True)
+            packed_key = "quant_state." + name
+            if packed_key in qs_dict:
+                return qs_dict[packed_key]
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
     def __getitem__(self, idx):
         """
         ensures compatibility with older quant state scheme with nested lists.

bitsandbytes/nn/modules.py

@@ -256,6 +256,43 @@ def __setstate__(self, state):
         self.bnb_quantized = state["bnb_quantized"]
         self.module = state["module"]
 
+    # Map from state_dict key names (as produced by QuantState.as_dict) to
+    # the actual QuantState attribute/access path. FSDP's _get_fqns() resolves
+    # dotted FQN keys via getattr, so "weight.quant_map" becomes
+    # getattr(weight, "quant_map") — we must map that to quant_state.code.
+    _QUANT_STATE_ATTR_MAP = {
+        # Direct QuantState attributes
+        "absmax": lambda qs: qs.absmax,
+        "code": lambda qs: qs.code,
+        "blocksize": lambda qs: qs.blocksize,
+        "dtype": lambda qs: qs.dtype,
+        "shape": lambda qs: qs.shape,
+        "offset": lambda qs: qs.offset,
+        "state2": lambda qs: qs.state2,
+        # as_dict serializes code → "quant_map"
+        "quant_map": lambda qs: qs.code,
+        "quant_type": lambda qs: qs.quant_type,
+        # as_dict serializes nested state2 attributes under "nested_*" keys
+        "nested_absmax": lambda qs: qs.state2.absmax,
+        "nested_blocksize": lambda qs: qs.state2.blocksize,
+        "nested_quant_map": lambda qs: qs.state2.code,
+        "nested_dtype": lambda qs: qs.state2.dtype,
+        "nested_offset": lambda qs: qs.offset,
+    }
+
+    def __getattr__(self, name):
+        # Proxy known QuantState attributes so that PyTorch's FSDP state_dict
+        # machinery (which traverses FQN paths via getattr) can find them.
+        accessor = self._QUANT_STATE_ATTR_MAP.get(name)
+        if accessor is not None:
+            quant_state = self.__dict__.get("quant_state")
+            if quant_state is not None:
+                try:
+                    return accessor(quant_state)
+                except AttributeError:
+                    pass
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
     def __deepcopy__(self, memo):
         new_instance = type(self).__new__(type(self))
         state = self.__getstate__()

tests/test_linear4bit.py

@@ -431,3 +431,69 @@ def test_linear4bit_torch_compile(device, quant_type, compute_dtype, compress_st
     grad_compiled = x.grad.clone()
 
     torch.testing.assert_close(grad_compiled, grad_ref)
+
+
+@pytest.mark.parametrize("device", get_available_devices())
+@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
+@pytest.mark.parametrize("compress_statistics", TRUE_FALSE, ids=id_formatter("compress_statistics"))
+def test_params4bit_quant_state_attr_access(device, quant_type, compress_statistics):
+    """Test that Params4bit proxies QuantState attributes for FSDP state_dict traversal (#1405).
+
+    PyTorch's FSDP state_dict machinery traverses FQN paths like
+    'model.layers.0.weight.absmax' using getattr(). This test verifies
+    that Params4bit and QuantState expose the attributes that appear as
+    state_dict keys so that _get_fqns() traversal succeeds.
+    """
+    if device == "hpu" and not is_supported_on_hpu(quant_type):
+        pytest.skip("This configuration is not supported on HPU.")
+
+    layer = bnb.nn.Linear4bit(
+        64,
+        64,
+        bias=False,
+        compress_statistics=compress_statistics,
+        quant_type=quant_type,
+    )
+    layer = layer.to(device)
+    w = layer.weight
+
+    assert w.quant_state is not None, "quant_state should be set after quantization"
+
+    # Direct QuantState attributes proxied through Params4bit
+    assert torch.equal(w.absmax, w.quant_state.absmax)
+    assert torch.equal(w.code, w.quant_state.code)
+
+    # "quant_map" is how as_dict() serializes "code" — FSDP uses this key name
+    assert torch.equal(w.quant_map, w.quant_state.code)
+
+    # QuantState packed key: as_dict(packed=True) produces "quant_state.bitsandbytes__<type>"
+    # FSDP resolves this as getattr(quant_state_obj, "bitsandbytes__<type>")
+    packed_attr = f"bitsandbytes__{quant_type}"
+    assert hasattr(w.quant_state, packed_attr)
+    packed_val = getattr(w.quant_state, packed_attr)
+    assert isinstance(packed_val, torch.Tensor)
+
+    # Simulate the full FSDP _get_fqns traversal for all state_dict keys
+    state_dict_keys = list(w.quant_state.as_dict(packed=True).keys())
+    for key in state_dict_keys:
+        # Each key is relative to "weight.", e.g. "absmax" or "quant_state.bitsandbytes__nf4"
+        parts = key.split(".")
+        obj = w
+        for part in parts:
+            obj = getattr(obj, part)
+        assert obj is not None
+
+    # hasattr should return True for proxied attrs, False for unknown ones
+    assert hasattr(w, "absmax")
+    assert hasattr(w, "code")
+    assert hasattr(w, "quant_map")
+    assert not hasattr(w, "nonexistent_attribute")
+
+    # Unknown attributes must still raise AttributeError
+    with pytest.raises(AttributeError, match="nonexistent_attribute"):
+        _ = w.nonexistent_attribute
+
+    # Verify that normal Params4bit attributes are unaffected by __getattr__
+    assert isinstance(w.quant_state, bnb.functional.QuantState)
+    assert isinstance(w.bnb_quantized, bool)
+    assert w.bnb_quantized is True