Fix pretrained checkpoint loading under torch < 2.4 (weights_only rejects hyper_parameters)#4
Open
alex-crlhmmr wants to merge 1 commit into
Conversation
load_checkpoint_state_dict guards add_safe_globals with hasattr, but that API does not exist before torch 2.4, so on older torch the guard is a no-op and torch.load(weights_only=True) rejects the argparse.Namespace that Lightning stores under hyper_parameters. Loading a pretraining checkpoint into the finetuning path then fails with an UnpicklingError. Fall back to weights_only=False when add_safe_globals is unavailable. The checkpoint is a local file produced by pretrain.py, so a full load is safe.
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
load_checkpoint_state_dictinmodels/checkpoint.pyguardstorch.serialization.add_safe_globalswithhasattr, but that API only exists from torch 2.4. On the torch version pinned byenvironment.yml(2.2.x) the guard is a silent no-op, andtorch.load(..., weights_only=True)then rejects theargparse.Namespacethat Lightning stores underhyper_parametersin every checkpoint produced bypretrain.py:So loading a pretraining checkpoint into the fine-tuning path (
--pretrain_checkpoint) fails on the project's own pinned environment.Change
Fall back to
weights_only=Falsewhenadd_safe_globalsis unavailable. The checkpoint is a local file produced bypretrain.py, so the full load is safe in that context. On torch >= 2.4 the behavior is unchanged (allowlistedweights_only=True).Verified
On torch 2.2.1 with a checkpoint from
pretrain.py: before the change,load_pretrained_encodersraises the UnpicklingError above; after the change, all three encoders load with zero missing and zero unexpected keys, and fine-tuning proceeds normally (we used this to fine-tune on FabWave and TMCAD).