Exposeweights_only for loading checkpoints with Trainer, LightningModule, LightningDataModule

[matsumotosan]

What does this PR do?

Fixes #20450 #20058 #20643

Before submitting

• Was this discussed/agreed via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

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

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

---

📚 Documentation preview 📚: https://pytorch-lightning--21072.org.readthedocs.build/en/21072/

[codecov]

Codecov Report

❌ Patch coverage is 55.55556% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 49%. Comparing base (634e6e6) to head (9aadbef).

❌ Your project check has failed because the head coverage (49%) is below the target coverage (99%). You can increase the head coverage or adjust the target coverage.

❗ There is a different number of reports uploaded between BASE (634e6e6) and HEAD (9aadbef). Click for more details.

HEAD has 70 uploads less than BASE

Flag	BASE (634e6e6)	HEAD (9aadbef)
cpu	41	27
lightning	24	16
python3.11	10	6
python	5	3
python3.10	10	6
python3.12.7	13	9
pytorch2.1	4	0
pytest-full	14	0
pytorch2.2.2	1	0
pytorch_lightning	6	0
pytorch2.7	2	0
pytorch2.3	3	0
pytorch2.6	2	0
pytorch2.5.1	1	0
pytorch2.4.1	1	0

Additional details and impacted files

@@            Coverage Diff            @@
##           master   #21072     +/-   ##
=========================================
- Coverage      87%      49%    -39%     
=========================================
  Files         269      266      -3     
  Lines       23520    23468     -52     
=========================================
- Hits        20503    11395   -9108     
- Misses       3017    12073   +9056     

[matsumotosan]

Should we default to weights_only=None or weights_only=True? If we have no use for weights_only=None, we can simplify the type hint to weights_only: bool = True.

[matsumotosan]

This is the simplest way that I could think of ensuring we continue testing the legacy checkpoints. Another way could be to use torch.serialization.add_safe_globals, but it seems a little more complicated (particularly since we're using the pl_legacy_patch context manager already.

[matsumotosan]

@Borda I wanted to get your opinion on something before moving forward.

I've added weights_only as an argument to LightningModule.load_from_checkpoint and all downstream functions to allow users to determine which option they want to use to load checkpoints.

My issue right now is with resuming training from a checkpoint with Trainer.fit. I see a few options right now:

1. Add weights_only as an argument to Trainer.fit (would also have to modify args for validate, test, and predict). Set default value to True.
2. Use weights_only=True everywhere, and print an error message advising user to set TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD if they would like to load with weights_only=False. Users must explicitly set environment variable to force loading with weights_only=False.
3. Add weights_only as an argument to Trainer initialization. Easy, but would not allow fine-grained control on loading models between different calls of fit, validate, etc.

I'm leaning towards option 1, but it involves changing up Trainer methods, which affects a lot of code so wanted to run this by you beforehand.

[Borda]

My issue right now is with resuming training from a checkpoint with Trainer.fit. I see a few options right now:

1. Add weights_only as an argument to Trainer.fit (would also have to modify args for validate, test, and predict). Set default value to True.
2. Use weights_only=True everywhere, and print an error message advising user to set TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD if they would like to load with weights_only=False. Users must explicitly set environment variable to force loading with weights_only=False.
3. Add weights_only as an argument to Trainer initialization. Easy, but would not allow fine-grained control on loading models between different calls of fit, validate, etc.

The cleanest way would probably be 1), but it brings so many new arguments for a marginal use... so personally I would go with 2)
cc: @lantiga

[lantiga]

Hi @matsumotosan let's do that only if the underlying torch is >= 2.6 (since starting weights_only became True by default from that point on), otherwise we're going to break a lot of older code

[matsumotosan]

I am not sure if it's possible to default to weights_only=True for trainer.{test,validate,test,predict}, since it loads a checkpoint at some point and that checkpoint may include elements that are not allowed by torch.load(..., weights_only=True) (Lightning itself saves classes like ModelCheckpoint), which cannot be loaded unless weights_only=False or a context manager is used.

The big issue with context managers is that a different one has to be used each time a different checkpoint is loaded. Setting the environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 does not work as it is overriden by the value passed to torch.load.

With this in mind, I think passing weights_only to trainer.{fit,validate,test,predict} may be the only short-term solution.

If we need to force the weights_only=True solution, every item saved in the checkpoint would need to be converted into a primitive type that torch.load(.., weights_only=True) would accept.

I have also added weights_only as an argument to the LightningModule and LightningDataModule classes as there are a few issues that point this out as a source of error:

• Expose weights_only option for loading checkpoints #20058
• convert_zero_checkpoint_to_fp32_state_dict fails with torch 2.6.0 due to weights_only default change #20643

Maybe we could default add weights_only and have it default to False for now until a future release so that we don't break users' code whilst adding an explicit option for this?

[Borda]

Suggested change

	if weights_only is None:
	if weights_only is None and _TORCH_GREATER_EQUAL_2_6:

[matsumotosan]

@Borda

torch.load_state_dict_from_url defaults to weights_only=False even for 2.8. Should we still default to weights_only=True in this case?

For cases of torch.load, I am not modifying weights_only if the function receives weights_only=None as the underlying torch.load function will handle defaults. We cannot do that with torch.load_state_dict_from_url as the argument is weights_only: bool = False.