Hook to transform the model before loading the weights

[awaelchli]

🚀 Feature

Motivation

This is a request from a user on Slack.
In their use case, they need to transform the model early in the trainer execution, after the checkpoint was loaded, but early enough before the weights get copied into the model (in case of executing trainer.fit(model, ckpt_path=...)). This is currently only done in a special case of the QuantizationAwareTraining callback.

Pitch

Provide a hook that runs before the model gets reloaded, but after the weights have been loaded from the file, i.e., in the sequence below the hook should run roughly where self._checkpoint_connector._restore_quantization_callbacks() runs.

https://github.com/Lightning-AI/lightning/blob/291267c3bff8054ec438960857c9f2fec1d54899/src/pytorch_lightning/trainer/trainer.py#L1071-L1079

The hook should take as input the checkpoint dict, so that the user can load their metadata.

def on_resume_start(self, lightning_module, trainer, checkpoint):
    """Do something with the model before restoration of the trainer/model state"""

Alternatives

One could make the QuantizationAwareTraining._load_before_model hook public, but this is limited to the use of the quantization callback only.

Additional context

Slack conversation

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Lite: enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.

• Bolts: Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.

• Lightning Transformers: Flexible interface for high-performance research using SOTA Transformers leveraging PyTorch Lightning, Transformers, and Hydra.

cc @Borda @tchaton @justusschock @awaelchli @carmocca @ananthsub @ninginthecloud @jjenniferdai @rohitgr7 @akihironitta

[rohitgr7]

might be useful with SWA too.

[d4l3k]

@awaelchli is there any timeline for something like this? I have a use case that does custom transforms on the model before loading the checkpoint state and exposing something like _load_before_model would be really nice to have

[awaelchli]

This isn't planned to be added at the moment unless there is a demand for it. Could you explain your use case, and maybe show some pseudo code, also explaining why applying your transformations can't be done in the existing hooks like LightningModule.load_state_dict or LightningModule.on_load_checkpoint()? That would be helpful :)

[d4l3k]

@awaelchli Similar to the old QuantizationAwareTraining callback we have some code that conditionally modifies the model using torch.fx graph transforms. Thus we need to be able to inspect and modify the model and state_dict before it gets loaded.

This transformation callback needs to be able to handle both transformed models (for failure recovery) as well as untransformed models (for fine tuning with transformations)

load_state_dict/on_load_checkpoint is too late since if the checkpoint was based on the transformed model the state_dict can't be loaded to the pre-transformed model.

In the removal PR for the QAT callback it says "You can copy the callback code and maintain it yourself." but that's impossible since the hook required doesn't exist. cc @lightningforever #16750

PyTorch native has some nn.Module hooks that could be used for this purpose as well but they're currently not public pytorch/pytorch#75287