container for quantization state components to work with Params4bit and similar classes
| 418 | |
| 419 | |
| 420 | class QuantState: |
| 421 | """container for quantization state components to work with Params4bit and similar classes""" |
| 422 | |
| 423 | valid_quant_types = ("fp4", "nf4") |
| 424 | valid_qs_type_keys = [f"bitsandbytes__{x}" for x in valid_quant_types] |
| 425 | valid_qs_keys = [ |
| 426 | "absmax", |
| 427 | "quant_map", |
| 428 | "nested_absmax", |
| 429 | "nested_quant_map", |
| 430 | "quant_state", |
| 431 | "quant_type", |
| 432 | "blocksize", |
| 433 | "dtype", |
| 434 | "shape", |
| 435 | "nested_blocksize", |
| 436 | "nested_dtype", |
| 437 | "nested_offset", |
| 438 | ] |
| 439 | |
| 440 | def __init__( |
| 441 | self, |
| 442 | absmax, |
| 443 | shape=None, |
| 444 | code=None, |
| 445 | blocksize=None, |
| 446 | quant_type=None, |
| 447 | dtype=None, |
| 448 | offset=None, |
| 449 | state2=None, |
| 450 | ): |
| 451 | self.absmax = absmax |
| 452 | self.shape = shape |
| 453 | self.code = code |
| 454 | self.dtype = dtype |
| 455 | self.blocksize = blocksize |
| 456 | self.quant_type = quant_type |
| 457 | self.offset = offset |
| 458 | self.state2 = state2 |
| 459 | self.nested = state2 is not None |
| 460 | |
| 461 | def __getattr__(self, name): |
| 462 | # Support attribute access for packed state_dict keys like "bitsandbytes__nf4". |
| 463 | # PyTorch's FSDP state_dict traversal (_get_fqns) resolves dotted FQN paths via |
| 464 | # getattr. The packed key "quant_state.bitsandbytes__nf4" causes it to call |
| 465 | # getattr(quant_state_obj, "bitsandbytes__nf4"), which we handle here. |
| 466 | if name.startswith("bitsandbytes__"): |
| 467 | qs_dict = self.as_dict(packed=True) |
| 468 | packed_key = "quant_state." + name |
| 469 | if packed_key in qs_dict: |
| 470 | return qs_dict[packed_key] |
| 471 | raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") |
| 472 | |
| 473 | def __getitem__(self, idx): |
| 474 | """ |
| 475 | ensures compatibility with older quant state scheme with nested lists. |
| 476 | assumes the following layout: |
| 477 | state = [qabsmax, input_shape, A.dtype, blocksize, [offset, state2], quant_type] |
no outgoing calls
no test coverage detected