Issues reloading model checkpoints in PL >= 1.6

[jdhorwood]

When using the Pretrained Goli fingerprint, I am getting the following error when reloading checkpoints with PL 1.6.1:

ModuleNotFoundError: No module named 'pytorch_lightning.utilities.argparse_utils'

It seems to be stemming from a backwards compatibility fix introduced in PL 1.5, but the module being deleted in the exit function of the patch is no longer there and I'm not sure why. Original PL patch here:

Lightning-AI/pytorch-lightning#9166

Note: I am able to load the fingerprint correctly, the issue arises after training with the fingerprint, and then loading a checkpoint from the final model. The issue only occurs for Goli fingerprint models.

[jdhorwood]

I think I figured out what is going on. The issue might be more relevant in molfeat so ping @maclandrol. I think the issue is being caused by a recursive call to with_legacy_patch() in pytorch lightning, which is messing around with sys.modules manually (probably a bad idea for a widely used framework).

After training a model using goli fingerprints, .load_from_checkpoint(best_ckpt_path) is called. This triggers a call to with_legacy_patch() which goes to its enter() function, which adds 'pytorch_lightning.utilities.argparse_utils' to sys.modules.

Then the load() function is called, which triggers a call to the Goli module's PredictorModule.load_from_checkpoint() within the featurizer. This then again calls with_legacy_patch() and enter().

As a result, the exit() function of the with_legacy_patch() block is called twice in a row. Since this function deletes 'pytorch_lightning.utilities.argparse_utils' from sys.modules in an effort to be self-consistent, an error is raised on the second call.

[maclandrol]

@JDavid04 this is indeed an issue with Pytorch lightning. I had it yesterday on unrelated context, and for some reasons even downgrading pl did not fix it (i didn't look too long, but there should be a pl-version that should be usable).

This is the fix I used: Basically the goal is to avoid going through load_from_checkpoint thus avoiding the with_legacy_patch context manager. If you are saving the hparams inside your plmodule, then this is perfectly fine:

MyPLModule = GreedyTuner # here for me. 
with fsspec.open(model_path, "rb") as IN:
    state = torch.load(IN)
    model = MyPLModule._load_model_state(state)

[jdhorwood]

Would this be possible within the load of the pretrained model? The issue only happens because of the nested calls to the context manager. I feel slightly more comfortable doing this fix on the lower level module rather than the higher level one.

[maclandrol]

I am a bit reluctant to patch the module like that for PL to fix the bug just after and we have to release again.

I will see, whether there is a context manager we can write that will patch the bug, but keep it more general.

[jdhorwood]

I think they won't patch it until an issue is raised because they added this is as a "feature" in v 1.5. So I will add an issue. I actually don't think we can fix it on our end because there is no way for us to inject the module in sys.modules between the inner context manager's calls to enter and exit I believe.

[maclandrol]

Let me try a few things, and I will ping you whether i was successful or not.

[maclandrol]

@JDavid04, I haven't forgotten this. I just did not like my solution.

[maclandrol]

Here (install mock first).

import sys
import pytorch_lightning

import contextlib
import mock
import pytorch_lightning.utilities.argparse
from types import ModuleType, TracebackType

class pl_legacy_patch_fix:    
    def __enter__(self) -> None:
        # `pl.utilities.argparse_utils` was renamed to `pl.utilities.argparse`
        legacy_argparse_module = ModuleType("pytorch_lightning.utilities.argparse_utils")
        sys.modules["pytorch_lightning.utilities.argparse_utils"] = legacy_argparse_module

        # `_gpus_arg_default` used to be imported from these locations
        legacy_argparse_module._gpus_arg_default = lambda x: x
        pytorch_lightning.utilities.argparse._gpus_arg_default = lambda x: x

    def __exit__(
        self, exc_type: type[BaseException] | None, exc_value: BaseException | None, exc_traceback: TracebackType | None
    ) -> None:
        if hasattr(pytorch_lightning.utilities.argparse, "_gpus_arg_default"):
            delattr(pytorch_lightning.utilities.argparse, "_gpus_arg_default")
        if "pytorch_lightning.utilities.argparse_utils" in sys.modules:
            del sys.modules["pytorch_lightning.utilities.argparse_utils"]

# Note where we are patching, not in the module defining pl_legacy_patch. 
#This is one of the gotcha of monkey patch
# I wish my other solution without redefining a contextmanager and just injecting the key into sys.modules would work 
with mock.patch("pytorch_lightning.core.saving.pl_legacy_patch", pl_legacy_patch_fix):
    model = MyPLModule.load_from_checkpoint(model_path)

[maclandrol]

You can ditch the mock dependency with:

import pytorch_lightning
@contextlib.contextmanager
def runtime_patch(object, name, patch):
    pre_patched_value = getattr(object, name)
    setattr(object, name, patch)
    yield object
    setattr(object, name, pre_patched_value)

with runtime_patch(pytorch_lightning.core.saving, "pl_legacy_patch", pl_legacy_patch_fix):
    model = GreedyTuner.load_from_checkpoint(model_path)

With some checks on pl version, this could be an "ok" alternative for now.

[jdhorwood]

Thanks Manu for checking this. I am wondering if this is a better solution though than the first one you proposed which bypasses the model loading? By the way I opened a PL ticket here : Lightning-AI/pytorch-lightning#12922

[maclandrol]

I also like the first one better and it's definitely more succinct. Unfortunately:

1. it's a bit annoying to have to use that everywhere. Because as you said the error occurs on >=2 calls to with_legacy_patch during a model loading. So it would still fails when the depth is >=3 and the other modules that are being loaded rely on load_from_checkpoint
2. You will have to deal with updating the hparams yourself on loading, which imo isn't an issue as I don't see anywhere in the codebase where we had to do that.

[jdhorwood]

What was the use case where it happened to you before? My feeling was that depth=2 would more or less be the max, and so doing this at the pretrained featurizer level would be sufficient for now. Wouldn't we also have to apply the patch above everywhere?

My feeling would be to probably restrict PL to < 1.5 until this is fixed on their side (which I am guessing would be relatively soon), but I am not sure if we now have requirements for newer PL versions with all the recent changes.

[maclandrol]

Chembag model for transfer, where the pretrained model has a plmodule as featuriser.

The second option should not have that issue.

[jdhorwood]

Will close this and revisit at a more appropriate location.