[fix][fsdp] Resume bitsandbytes checkpoints strictly

[bvolpato]

Summary

FSDP checkpoints for bitsandbytes 4-bit modules include quantization metadata beside each weight. Freshly constructed modules can accept the weight while reporting that metadata as unexpected, causing strict resume to fail before optimizer, scheduler, RNG, or dataloader state is restored.

Changes

• Collect model-load incompatibilities without immediately raising.
• Accept known FP4 and NF4 metadata only when its associated base weight was accepted.
• Preserve strict failure for missing keys and unrelated unexpected keys.
• Leave non-strict loading unchanged.

Checkpoint format and model-loading behavior outside bitsandbytes metadata remain unchanged.

Validation

• Focused checkpoint tests (9 passed)
• Ruff, Black, Gitleaks, and git diff --check
• Single-GPU NF4 QLoRA resume from step 20 to 40, then step 40 to 60, with trainer, dataloader, policy, optimizer, scheduler, and RNG state restored
• Nonzero gradients observed and checkpoints saved after both resumed phases

[gemini-code-assist]

Code Review

This pull request introduces a custom model state dict loading helper, _load_model_state_dict, to support strict loading while ignoring bitsandbytes quantization metadata keys (such as .absmax, .quant_map, etc.). It also adds corresponding unit tests to verify this behavior. The review feedback suggests optimizing the quantization key check by leveraging str.endswith with a tuple of suffixes to avoid unnecessary loops, and simplifying the calculation of loaded keys using set operations on dictionary keys.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.

^{Reviewed by Cursor Bugbot for commit e283a48. Configure here.}

[cursor]

Strict load crashes on key set math

High Severity

loaded_keys is computed with state_dict.keys() - incompatible.unexpected_keys. PyTorch returns unexpected_keys as a list, and dict_keys subtraction only accepts a set, so strict resume raises TypeError before any bitsandbytes filtering runs.

 

^{Reviewed by Cursor Bugbot for commit e283a48. Configure here.}

[bvolpato]

This is not reproducible on supported Python. Both dict_keys and odict_keys accept list operands for subtraction and return a set: {'a': 1, 'b': 2}.keys() - ['b'] == {'a'}. Focused strict-load tests also pass with PyTorch's list-valued unexpected_keys.

[cursor]

Nested quant suffixes match too early

High Severity

_is_bnb_quantization_state_key now returns on the first suffix that matches by endswith, so .nested_absmax and .nested_quant_map are stripped with .absmax / .quant_map instead. Valid double-quant metadata is then treated as unexpected and strict resume fails.

 

^{Reviewed by Cursor Bugbot for commit e283a48. Configure here.}

[bvolpato]

These suffixes do not overlap because their delimiters differ: weight.nested_absmax.endswith(.absmax) and weight.nested_quant_map.endswith(.quant_map) are both false. Focused tests cover all six supported suffixes, including both nested forms, and strict loading passes.