unpacks components of state_dict into QuantState where necessary, convert into strings, torch.dtype, ints, etc. qs_dict: based on state_dict, with only relevant keys, striped of prefixes. item with key `quant_state.bitsandbytes__[nf4/fp4]` may contain minor and non
(cls, qs_dict: dict[str, Any], device: torch.device)
| 492 | |
| 493 | @classmethod |
| 494 | def from_dict(cls, qs_dict: dict[str, Any], device: torch.device) -> "QuantState": |
| 495 | """ |
| 496 | unpacks components of state_dict into QuantState |
| 497 | where necessary, convert into strings, torch.dtype, ints, etc. |
| 498 | |
| 499 | qs_dict: based on state_dict, with only relevant keys, striped of prefixes. |
| 500 | |
| 501 | item with key `quant_state.bitsandbytes__[nf4/fp4]` may contain minor and non-tensor quant state items. |
| 502 | """ |
| 503 | |
| 504 | # unpacking tensor with non-tensor components |
| 505 | qs_key = [k for k, v in qs_dict.items() if "quant_state" in k and isinstance(v, torch.Tensor)] |
| 506 | if "quant_type" not in qs_dict: |
| 507 | if not qs_key: |
| 508 | raise ValueError("Expected packed or unpacked quant_state items, found neither") |
| 509 | elif len(qs_key) != 1 or qs_key[0].split(".")[-1] not in cls.valid_qs_type_keys: |
| 510 | raise ValueError( |
| 511 | f"There should be exactly one `quant_state` item with ending from {cls.valid_qs_type_keys}.\nDetected {qs_key}.", |
| 512 | ) |
| 513 | |
| 514 | # unpacking minor and non-tensor quant state items if necessary |
| 515 | if len(qs_key) == 1: |
| 516 | first_qs_key = qs_key[0] |
| 517 | qs_dict.update(unpack_tensor_to_dict(qs_dict.pop(first_qs_key))) |
| 518 | |
| 519 | qs_dict = {k.split(".")[-1]: v for k, v in qs_dict.items()} # strip prefixes |
| 520 | assert set(qs_dict.keys()).issubset(cls.valid_qs_keys) |
| 521 | |
| 522 | if "nested_absmax" in qs_dict: |
| 523 | offset = torch.tensor(float(qs_dict["nested_offset"])).to(device) |
| 524 | state2 = cls( |
| 525 | absmax=qs_dict["nested_absmax"].to(device), |
| 526 | blocksize=qs_dict["nested_blocksize"], |
| 527 | code=qs_dict["nested_quant_map"].to(device), |
| 528 | dtype=getattr(torch, qs_dict["nested_dtype"]), |
| 529 | ) |
| 530 | else: |
| 531 | offset, state2 = None, None |
| 532 | |
| 533 | quant_state = cls( |
| 534 | quant_type=qs_dict["quant_type"], |
| 535 | absmax=qs_dict["absmax"].to(device), |
| 536 | blocksize=qs_dict["blocksize"], |
| 537 | code=qs_dict["quant_map"].to(device), |
| 538 | dtype=getattr(torch, qs_dict["dtype"]), |
| 539 | shape=torch.Size(qs_dict["shape"]) if qs_dict["shape"] is not None else None, |
| 540 | offset=offset, |
| 541 | state2=state2, |
| 542 | ) |
| 543 | return quant_state |
| 544 | |
| 545 | def as_dict(self, packed: bool = False) -> dict[str, Any]: |
| 546 | """ |