Backward compatibility for corrdiff checkpoints

[CharlelieLrt]

PhysicsNeMo Pull Request

Description

This PR introduces a patch to be able to load existing UNet and SongUNet checkpoints, similar to those used in CorrDiff.

To apply backward compatibility patches and be able to load old checkpoints, run:

model = Module.from_checkpoint(chckpt.mdlus, backward_compatibility=True)

Bu default backward_compatibility=False, in which case the checkpoint loader will raise an error when attempting to load an old checkpoint.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.

Dependencies

None.

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

@pzharrington
Following our discussion offline, I attempted a refactor based on a function to explicitly convert checkpoints such as convert_checkpoint(old_chckp.mdlus, new_chckpt.mdlus). But I realized that this approach might not work for 2 reasons:

1. In utils_compatibility.py I need a statement if cls is ... to identify the class for which we are applying the patch. This means that this can only be done after cls has been determined, so within the instantiate. This cannot be done in a separate convert function.
2. If going for an explicit convert function, we will need to update all our affected checkpoints hosted on NGC (we can't really expect users to always run convert for those). But updating multiple checkpoints on NGC could quickly become unmanageable if more checkpoints are affected.

So, I proposed another refactor that does essentially the same thing, but that is enabled by a keyword argument passed to load_checkpoint: model = Module.from_checkpoint(chckpt.mdlus, backward_compatibility=True).

[pzharrington]

I think the keyword arg to toggle backward compat mode in from_checkpoint provides a reasonable degree of safety, having users need to opt into loading an older checkpoint which is good.