Move code from TrainerOptimizersMixin into TrainingTypePlugin

[daniellepintz]

Proposed refactor

Move all the code from this class to TTP:
https://github.com/PyTorchLightning/pytorch-lightning/blob/2036dfb5df2a48056d1c043ee311aed793187070/pytorch_lightning/trainer/optimizers.py#L28

Motivation

After #10596 the TTP will own the optimizer related logic previously owned by Accelerator. It makes sense for the TTP to also own the optimizer related code in TrainerOptimizersMixin. This would allow us to avoid this type of circuitous logic: https://github.com/PyTorchLightning/pytorch-lightning/blob/dcafc95f2b0fd3f176d425139ca99676ce943a12/pytorch_lightning/plugins/training_type/training_type_plugin.py#L244-L245

We in general want to reduce the number of mixins (#10417)

cc @justusschock @awaelchli @akihironitta

[awaelchli]

Thanks for the initiative.

I propose to only move the contents of the init_optimizers method to the TTP. In there, we would check if the model has a configure_optimizers method and call it.
The static methods could become functions and be part of core (LightningModule) because they are responsible for parsing the output of the LightningModule hook. They don't belong into the TTP imo.

[daniellepintz]

@awaelchli Sounds good!!

[daniellepintz]

and be part of core (LightningModule)

@awaelchli do you mean in lightning.py?

[awaelchli]

Yes, basically in core/lightning.py or core/optimizer.py or somewhere there where it's close to the hook that retrieves the optimizers and schedulers from the user (configure_optimizers).

[ananthsub]

@awaelchli is something like this what you had in mind?

class LightningModule(...):
    
    @abstractmethod
    def configure_optimizers(self):
        ....

    def init_optimizers_and_lr_schedulers(self) -> TypedDict[...]:
        output = self.configure_optimizers()
        # parse & validate output
        # return components

then we could also schematize/add typing for configure_optimizers so all the typechecking & validation is in one place, with the strategy calling init_optimizers as opposed to configure_optimizers

[daniellepintz]

@ananthsub it sounds like what you are saying is to move init_optimizers to the LM as well, rather than putting it in TTP?

[ananthsub]

@ananthsub it sounds like what you are saying is to move init_optimizers to the LM as well, rather than putting it in TTP?

I think the caller of configure_optimizers should be the strategy. However, the definition of the parsing & validation logic from the output of configure_optimizers should be shared.