[Afloat16]

Problem

Bitsandbytes optimizers inherit register_load_state_dict_pre_hook and register_load_state_dict_post_hook, but their load_state_dict override never calls them. Checkpoint adaptation therefore silently fails: a registered pre-hook cannot migrate saved parameter groups before validation, and a registered post-hook cannot run after restoration.

Change

Invoke the existing pre-hook registry on a shallow copy, honoring replacement dictionaries, then retain the existing deep copy, FSDP quantization-state unwrapping, dtype/device handling and validation. Invoke post-hooks only after the state has been installed successfully. Ordering, prepend=True, and handle removal use the inherited registries.

The loader continues to avoid modifying checkpoint data during its own unwrap/cast operations. User pre-hooks keep PyTorch's normal shallow-copy semantics for nested data.

Validation

Base a616886cacf86701514f0cb1d6be760f967481c0; patch 5fc14afef96ad3e2c2bfd387b8b3587b861082cf.

• The same 11 new cases: 11 failed before; 11 passed after.
• New cases plus existing AdEMAMix checkpoint round trips: 15 passed, 248 deselected.
• Existing non-paged FP32 optimizer tests across four parameter shapes: 28 passed, 235 deselected.
• Real native CPU Adam8bit, Adam32bit, SGD8bit and SGD32bit train a two-layer network for three steps, serialize a checkpoint, migrate one saved group into two groups through a pre-hook, and resume for three more steps. Model outputs, parameters and optimizer state match uninterrupted training exactly.
• Covers in-place and replacement pre-hooks, post-hook observation of the installed state, ordering, handle removal, move_to_device flags, and validation failure without a post-hook call.
• 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_optim.py -k 'load_state_dict_hook or load_state_dict_post_hook or ademamix_state_dict' -q

Python 3.12.14 / PyTorch 2.15.0+cpu; current bitsandbytes native CPU library built from source. GPU/distributed execution and the complete test suite were not run.

Checked the optimizer file history, current PR catalog and all-state checkpoint/hook searches. The related FSDP prototype #2081 changes an example, and #1958 fixes module serialization; neither implements optimizer load hooks.