Expose weights_only argument in Fabric.load #21470
Open
+34
−9
Conversation
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
CodeVishal-17
requested review from
ethanwharris,
justusschock,
lantiga and
tchaton
as code owners
for more information, see https://pre-commit.ci
GdoongMathew
suggested changes
Jan 6, 2026
src/lightning/fabric/fabric.py
Outdated
Comment on lines
877
to
882
| """Load a checkpoint from a file and restore the state of objects (modules, optimizers, etc.) How and which | ||
| processes load gets determined by the `strategy`. | ||
| How and which processes load gets determined by the `strategy`. | ||
| This method must be called on all processes! | ||
| Args: | ||
| path: A path to where the file is located. | ||
| state: A dictionary of objects whose state will be restored in-place from the checkpoint path. | ||
| If no state is given, then the checkpoint will be returned in full. | ||
| strict: Whether to enforce that the keys in `state` match the keys in the checkpoint. | ||
| Returns: | ||
| The remaining items that were not restored into the given state dictionary. If no state dictionary is | ||
| given, the full checkpoint will be returned. | ||
| Example:: | ||
| # Load full checkpoint | ||
| checkpoint = fabric.load("checkpoint.pth") | ||
| # Load into existing objects | ||
| state = {"model": model, "optimizer": optimizer} | ||
| remainder = fabric.load("checkpoint.pth", state) | ||
| epoch = remainder.get("epoch", 0) | ||
| """ |
Contributor
There was a problem hiding this comment.
The original docstring should be extended with the additional kwarg weight_only instead of removed.
Author
There was a problem hiding this comment.
I’ve updated the docstring to preserve the original content and extended it with the weights_only argument, including fixing the example formatting. Thanks for the review!
Contributor
…hal-17/pytorch-lightning into fabric-load-weights-only
for more information, see https://pre-commit.ci
GdoongMathew
reviewed
Jan 6, 2026
Comment on lines
905
to
906
| # We need to unwrap objects (see above) but this creates a new dictionary. In-place updates | ||
| # (for user metadata) wouldn't show up in the original dict, so we need to copy the data back. |
Contributor
There was a problem hiding this comment.
The comments here should still be preserved.
Author
There was a problem hiding this comment.
Thanks for the review 👍
I’ve resolved the merge conflict in Fabric.load, preserved all existing comments and examples, and kept the
weights_only change purely additive. Please let me know if everything looks good now.
…/CodeVishal-17/pytorch-lightning into fabric-load-weights-only" This reverts commit 6143e65, reversing changes made to debde18.
…hal-17/pytorch-lightning into fabric-load-weights-only
for more information, see https://pre-commit.ci
GdoongMathew
approved these changes
Jan 7, 2026
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #21470 +/- ##
=======================================
- Coverage 87% 87% -0%
=======================================
Files 270 270
Lines 24059 24059
=======================================
- Hits 20862 20859 -3
- Misses 3197 3200 +3 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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?
Exposes the
weights_onlyargument inFabric.load()and forwards it to the underlying strategy checkpoint loading call.This restores compatibility with mixed-object checkpoints after the PyTorch 2.6 change where
torch.loaddefaults toweights_only=True.Motivation and context
With PyTorch 2.6,
torch.loaddefaults toweights_only=True, which prevents loading checkpoints that contain non-tensor objects (such as optimizers, schedulers, or metadata).While
strategy.load_checkpoint(..., weights_only=False)already supports this use case,Fabric.load()did not expose the argument, making it impossible for users to opt out of the new default behavior when using Fabric.This PR adds a keyword-only
weights_onlyargument toFabric.load()without changing the existing default behavior.Fixes #21459
Before submitting
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
📚 Documentation preview 📚: https://pytorch-lightning--21470.org.readthedocs.build/en/21470/