Missing / Unexpected key(s) when loading a checkpoint from a compiled model

[konstantinjdobler]

Bug description

When loading a checkpoint that was compiled using torch.compile on a using

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

[awaelchli]

@konstantinjdobler We have only tested Lightning with compiling the LightningModule directly. That is probably what most people want. But of course, your use case of compiling only a submodule is also valid. The problem is that torch.compile returns a new type OptimizedModule that wraps around the original module. This is why an extra prefix appears in the checkpoint.

To fix this, we would need to iterate through all submodules and unwrap the modules of type OptimizedModule when saving a checkpoint through .state_dict().

[konstantinjdobler]

@awaelchli That's kind of what I thought. In my case, I am actually fine compiling the entire LightningModule but I can imagine cases where individual submodules cannot be compiled because of custom operations etc. but one might want to compile the other submodules.

[RahulSajnani]

A quick fix would be to initialize and compile model within configure_model(). Ensure you initialize the model and then compile in the same function as it didn't work otherwise for me