Skip to content

[fix][fsdp] Resume bitsandbytes checkpoints strictly - #2007

Draft
bvolpato wants to merge 1 commit into
NovaSky-AI:mainfrom
bvolpato:bvolpato/fsdp-bnb-checkpoint-resume
Draft

[fix][fsdp] Resume bitsandbytes checkpoints strictly#2007
bvolpato wants to merge 1 commit into
NovaSky-AI:mainfrom
bvolpato:bvolpato/fsdp-bnb-checkpoint-resume

Conversation

@bvolpato

@bvolpato bvolpato commented Aug 10, 2026

Copy link
Copy Markdown
Contributor

Problem

FSDP checkpoints for bitsandbytes 4-bit modules contain auxiliary quantization state next to each weight, including absmax, quantization maps, nested maps, and serialized FP4 or NF4 state. A freshly constructed module can consume its weight while reporting those auxiliary entries as unexpected. SkyRL currently requests strict loading, so resume stops before optimizer, scheduler, RNG, or dataloader state can be restored.

Change

  • collect load incompatibilities without immediately raising
  • accept known bitsandbytes metadata only when associated base weight was accepted by model loader
  • preserve strict failure for every missing key and unrelated unexpected key
  • leave non-strict loading unchanged

Why this belongs in SkyRL

SkyRL writes and restores FSDP checkpoint state. Resume must understand metadata emitted by supported 4-bit training modules. Fix is checkpoint-format compatibility, independent of model architecture or rollout backend. It complements QLoRA support in #1968 without combining model-loading changes into this PR.

Compatibility and safety

  • checkpoint format does not change
  • FP4 and NF4 metadata are supported
  • suffix match alone is insufficient: metadata must map to checkpoint weight accepted by module
  • unrelated keys such as unexpected.absmax still fail strict loading
  • missing parameters still fail strict loading
  • compatibility check reuses existing load result and does not traverse model state again

Validation

  • focused checkpoint tests: 4 passed
  • resumed single-GPU NF4 QLoRA checkpoint from step 20
  • restored trainer, dataloader, policy, optimizer, scheduler, and RNG state
  • continued to step 40 with nonzero gradients and saved checkpoint
  • resumed step 40, continued to step 60, and saved another checkpoint
  • changed-file Ruff, Black, and hardcoded-secret checks pass

CI note

Current full-tree Black failure is already present on main in tests/tinker/test_engine.py and is isolated in #2009.

@bvolpato
bvolpato force-pushed the bvolpato/fsdp-bnb-checkpoint-resume branch from fcdd6d9 to 52951bc Compare August 10, 2026 00:56
@bvolpato bvolpato changed the title [fix][fsdp] Resume bitsandbytes checkpoints [fix][fsdp] Resume bitsandbytes checkpoints strictly Aug 10, 2026
@bvolpato
bvolpato force-pushed the bvolpato/fsdp-bnb-checkpoint-resume branch from 52951bc to 035c703 Compare August 10, 2026 01:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant