Customizing model checkpointing logic

[ZhaofengWu]

I don't think overriding ModelCheckpoing.save_checkpoint() is the best idea, since it contains logic such as finding topk, etc. I then traced down the code from ModelCheckpointCallback, and eventually ended up at TrainingTypePlugin, and specifically its lightning_module_state_dict() method which seems like what I should inherit, overload, and pass to the Trainer's constructor as a plugin. However, which class should I inherit from? The generic TrainingTypePlugin? A specific subclass such as SingleDevicePlugin? Neither sound general enough.

On the other hand, I'm inclined to directly override state_dict(), but I worry that it might be used for purposes other than checkpointing and might cause silent bugs.

This might be related to #9627.

[ananthsub]

Could you describe what functionality you'd like to add during checkpointing? This would help us determine a recommendation for which component to look at extending

[ZhaofengWu]

I would like to choose only a few parameters to save, but not all. But in general I think it'd be great if a save_checkpoint-like method could be easily accessed and overridden to arbitrarily change the saving logic, for example as a LightningModule method.

[ananthsub]

I think it'd be great if a save_checkpoint-like method could be easily accessed and overridden to arbitrarily change the saving logic, for example as a LightningModule method.

It already is. The LightningModule extends nn.Module, so you can override state_dict to return only the parameters you care about.

Or if you want to mutate the entire checkpoint before saving, you can do so in LightningModule.on_save_checkpoint

[ZhaofengWu]

I don't believe that's the right the thing to do: in the pytorch documentation, state_dict() is described as "a dictionary containing a whole state of the module." So, the other parameters are still technically a part of the "whole state of the module", but I just don't need to persist them in a checkpoint. In other words, some other piece of code might call state_dict() to inspect the "state of the module" for purposes other than checkpointing, and it would be wrong to only return parameters that I'd like to persist. Therefore, I think there should be a similar method only for the purpose of checkpointing.

on_save_checkpoint seems to serve my purpose, though. Thanks for the pointer! Conceptually it might be simpler for me to customize what to return directly, rather than having lightning generate a state dict first and then me modifying it. Nevertheless in practice it probably doesn't make a difference. I still do think a save_checkpoint-like method is better, since it can enable more customization. For example, it allows skipping the checkpoint by returning without doing anything, but on_save_checkpoint can't. Granted, there are other mechanisms in lightning that achieves that specific purpose, but I think more flexibility here is better and I believe more aligned to lightning's philosophy.