Fix QuantState.as_dict

Signed-off-by: cyy <[email protected]>

Author: cyyever

bitsandbytes/functional.py

@@ -503,7 +503,7 @@ def from_dict(cls, qs_dict: dict[str, Any], device: torch.device) -> "QuantState
         )
         return quant_state
 
-    def as_dict(self, packed=False):
+    def as_dict(self, packed: bool = False) -> dict[str, Any]:
         """
         returns dict of tensors and strings to use in serialization via _save_to_state_dict()
         param: packed -- returns dict[str, torch.Tensor] for state_dict fit for safetensors saving
@@ -532,7 +532,10 @@ def as_dict(self, packed=False):
         # packed format allows serialization of non-tensor components, critical for saving in safetensors format
         qs_packed_dict = {k: v for k, v in qs_dict.items() if isinstance(v, torch.Tensor)}
         non_tensor_dict = {k: v for k, v in qs_dict.items() if not isinstance(v, torch.Tensor)}
-        qs_packed_dict["quant_state." + "bitsandbytes__" + self.quant_type] = pack_dict_to_tensor(non_tensor_dict)
+        key = "quant_state.bitsandbytes__"
+        if self.quant_type is not None:
+            key += self.quant_type
+        qs_packed_dict[key] = pack_dict_to_tensor(non_tensor_dict)
         return qs_packed_dict
 
     def to(self, device):