Lightning-AI / Lightning-AI/pytorch-lightning
Missing / Unexpected key(s) when loading a checkpoint from a compiled model
@awaelchli is already working on this.
Since Dec 29, 2023.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
When loading a checkpoint that was compiled using `torch.compile` on a using
```python
model = LitModel.load_from_checkpoint(...)
...
#later on
model.submodule = torch.compile(model.submodule)
```
I am getting missing key(s) and unexpected key(s) errors. Specifically, all the model keys are "missing" and instead there are "unexpected" keys where the compiled module is replaced by `_orig_mod`. Is it just not possible to load a checkpoint from a checkpoint containing a compiled module? Going into the checkpoint and replacing the keys by the correct value is not hard, but inconvenient.
I have not tested whether compiling the entire LightningModule works (perhaps it is just a legitimate limitation that lightning cannot know whether a submodule was compiled or not).
### What version are you seeing the problem on?
`lightning==2.0.0`
cc @borda @carmocca
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.