Expose `weights_only` option for loading checkpoints

[awaelchli]

Description & Motivation

PyTorch will default the torch.load(weights_only=) option to True in the future, and started emitting FutureWarning in 2.4. Since Lightning allows storing all kinds of data into the checkpoint (mainly through save_hyperparameters() and the on_save_checkpoint hooks), so far we have to explicitly set weights_only=False internally to be able to load checkpoints. But we could also expose this argument as a flag for the user to choose how strict they want to be when loading checkpoints.

Pitch

In Fabric.load, expose weights_only.
In PyTorch Lightning, either add fit(..., load_weights_only=...) or add an attribute to the LightningModule (similar to how we handle strict_loading).

Alternatives

No response

Additional context

No response

cc @Borda @awaelchli

[ryan597]

Is this issue closed with #20010 merged? @awaelchli

[awaelchli]

No it isn't, because in #20010 we just set weights_only=False (because we have to). But a user wouldn't be able to control it (unless they set the environment variable).

Why do you ask @ryan597? :)

[ryan597]

Ah okay, I was looking at some issues to pickup :)

[awaelchli]

Cool. This one here probably needs a bit validation first that users actually need it.
If you're interested in picking up issues to work on, that's great and appreciated. Feel free to contact me on slack or discord and I can point you to some based on your interest. The ones labelled with help wanted are generally ones that are open to work on.

[yair-schiff]

Any update on this? torch 2.6 changed default behavior to have weights_only=True which is breaking my workflow of trainer.fit(ckpt_path=ckpt_path)

[katosad]

Any update on this? My org is using torch 2.6.0 and a core inference flow is now broken as I have no way to specify weights_only to load_from_checkpoint function. Would love if y'all could get this implemented. Please let me know if you have any workaround suggestions.

[yair-schiff]

@katosad, a bit hacky, but I think you can set TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 environment variable to avoid this breaking behavior for now. Hope that helps!

(ref: pytorch docs)