[Afloat16]

Problem

A quantized Linear8bitLt moves its weight and scale caches into self.state on the first forward. Its checkpoint loader subsequently looks only at weight.SCB, which is now None, so reloading an otherwise valid checkpoint incorrectly raises the error for an unquantized module. This also prevents switching between saved quantized weights after inference or training has started.

Change

Restore scales into the active scale buffer, using state.SCB when the weight-owned buffer has already moved. Keep the state reference bound to that buffer after copying checkpoint values. The existing error for genuinely uninitialized quantization buffers remains in place.

Validation

Base: a616886cacf86701514f0cb1d6be760f967481c0.

• Identical 12 new native cases: 12 failed before; 12 passed after.
• The complete tests/test_linear8bitlt.py module excluding compiler-specific cases: 47 passed, 4 skipped, 16 deselected.
• Covers checkpoints saved before/after first forward, cold/training/evaluation destinations, and outlier thresholds 0 and 6.
• Every case serializes and sequentially restores two different quantized checkpoints, checks exact outputs, scale values, weight/cache storage identity, unchanged caller checkpoint data and resaved state.
• Three native SGD bias-training steps after each restoration compare outputs, input gradients, bias gradients and updated biases exactly against the source model.
• Repository-pinned Ruff 0.14.3 checks/formatting and git diff --check pass.

BNB_TEST_DEVICE=cpu OMP_NUM_THREADS=2 python -m pytest tests/test_linear8bitlt.py -k 'not torch_compile' -q

Python 3.12.14 / PyTorch 2.15.0+cpu, with the current native CPU library built from source. GPU execution, compiler-specific tests and the full repository suite were not run.

The failure described in #960 reproduces on this base. Checked all-state checkpoint/reload searches and the actual related serialization changes #159, #1099, #1865 and #2129; none restores the state-owned scale buffer on this load path.