Repository navigation
Conversation
`FSDPStrategy.load_checkpoint(..., strict=False)` (Fabric and Trainer) and the Fabric ModelParallel strategy only forwarded `strict` to `module.load_state_dict`, after `torch.distributed.checkpoint.load` had already raised `Missing key in checkpoint state_dict` for any model key absent from the sharded checkpoint. Pass a `DefaultLoadPlanner(allow_partial_load=not strict)` to the distributed checkpoint load so that missing keys are tolerated when `strict=False`, while `strict=True` keeps raising. Optimizer state loading is unchanged. Fixes Lightning-AI#20274 Signed-off-by: Parth Vadera <parthvadera@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
deep1blue
requested review from
ethanwharris,
justusschock and
tchaton
as code owners
October 19, 2026 00:24
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes #20274
FSDPStrategy.load_checkpoint(..., strict=False)(in both Fabric and Trainer) and Fabric'sModelParallelStrategy.load_checkpoint(..., strict=False)previously only forwardedstricttomodule.load_state_dict. However, for sharded (distributed) checkpoints,_distributed_checkpoint_loadinvokedtorch.distributed.checkpoint.load(module_state, checkpoint_id=path)with the defaultDefaultLoadPlanner(allow_partial_load=False), which raisedCheckpointException(RuntimeError: Missing key in checkpoint state_dict) beforemodule.load_state_dict(..., strict=False)was ever reached.Changes
strict: bool = Trueto_distributed_checkpoint_loadand passplanner=DefaultLoadPlanner(allow_partial_load=not strict)totorch.distributed.checkpoint.load.strict=strictfromFabricFSDPStrategy.load_checkpointandModelParallelStrategy.load_checkpoint, andstrict=self.lightning_module.strict_loadingfromTrainerFSDPStrategy.load_checkpointwhen loading module state.test_distributed_checkpoint_load_strictintests/tests_fabric/strategies/test_fsdp.pyand update existing_distributed_checkpoint_loadmocks intests/tests_fabric/strategies/test_model_parallel.py.PR review
Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:
Reviewer checklist
Credits & Acknowledgments
strict = Falsedoes not work when the checkpoint is distributed #20274 and suggestingDefaultLoadPlannerin the issue report, and to @nocoding03 for proposingallow_partial_load=not strictand threadingstrictinto_distributed_checkpoint_load.