PyTorch Lightning & Jupyter Compatibility

[Feuermagier]

This PR introduces three changes:

• Find the device of each parameter group on demand, instead of at initialization time. This is needed for compatibility with PyTorch lightning, which moves the model to the device only after the optimizer has been created. It also allows to use a different device per parameter group.
• Evaluate the step closure inside of the with self.sampled_params(train=True) context manager. PT lightning requires this as it uses the closure. This change also brings the implementation more in line with the PT documentation, which says that the closure should only evaluate the model and compute gradients.
• Wrap the yield in the sampled_params generator with try-finally, so that the parameter average gets restored even if the parent function caused an exception. This useful for Jupyter notebooks, where crashing cells can otherwise leave the optimizer in an inconsistent state.

I've been using this fork for the past few weeks, so it's at least somewhat tested. Please let me know if you'd like me to change anything.

[Leon0402]

@Feuermagier Ah pretty cool! Just saw this after I posted a very similar issue #5 actually.
For 1. I think configure_optimizers is called when the model is already on the gpu, at least I did not have any issues with this. But it might be a different story for validation, testing, checkpoint loading (I had such issues there)
For 2. there is also a somewhat easy workaround in the issue, but the changes here are nicer of course.

As you seem to have worked with PL and IVON, I was wondering whether you have some public repo with the code. Beside the mentioned points, I struggled a bit with validation / testing and checkpoint loading as PL makes the assumption that optimizer is only needed for training. I found a somewhat hacky solution for this:

    def on_validation_start(self) -> None:
        self._load_optimizer()

    def on_test_start(self) -> None:
        self._load_optimizer()

    @override
    def configure_optimizers(self):
        self.optimizer = self._load_optimizer()
        return {
            "optimizer": self.optimizer,
            "lr_scheduler": {
                "scheduler": self._configure_scheduler(self.optimizer),
                "interval": "step",
            },
        }
        
     @override
    def on_save_checkpoint(self, checkpoint: dict[str, Any]):
        super().on_save_checkpoint(checkpoint)
        if self.trainer.is_global_zero:
            torch.save(
                self.optimizer.state_dict(),
                Path(self.trainer.checkpoint_callback.dirpath) / "optimizer.pth",
            )

    @override
    def on_load_checkpoint(self, checkpoint: dict[str, Any]):
        super().on_load_checkpoint(checkpoint)
        if self.ivon_parameters.optimizer_state_path:
            logging.warning("Overriding optimizer state path")
        self.ivon_parameters.optimizer_state_path = Path(self.trainer.ckpt_path).parent / "optimizer.pth"

So essentially savin IVON explicitly in a file and loading it on validation / testing. It works, but it is a bit overhead (and it does not work nicely with multiple Checkpoint Callbacks for instance). So I was wondering whether you have a more elegant solution :-)

Also tagging some of the repo owerns @cong-bai @moellenh @pnickl as this MR has been open quite a while and was probably overlooked. Would be great if someone could officially review and merge this!

[Feuermagier]

I'm actually using my own validation/testing infrastructure, not PL's. But another solution would be to make IVON's state part of the model state (though that might increase the size of the checkpoints during training since the IVON state gets saved two times):

class ModuleWithIVON(nn.Module):
    def __init__(self):
        super().__init__():
        self.model = # ...
        self.ivon = IVON(self.model.parameters(), ...)

    def get_extra_state(self):
        return {
            "ivon": self.ivon.state_dict(),
        }

    def set_extra_state(self, state):
        self.ivon.load_state_dict(state["ivon"])

and then overwrite the ivon attribute in configure_optimizers() (or just return self.model.optimizer). I don't think this is more elegant than your solution though...