Repository: pytorch
Issue: `LRScheduler` Refactoring
Labels: module: optimizer,triaged,better-engineering,module: LrScheduler

### 🚀 The feature, motivation and pitch

## Feature
Proposal to solve issue mainly raised from issue #67760, and some other issues related to `lr_scheduler` module.

* Refactor `lr_scheduler`  module
  * make `lr_scheduler` class to support all hyperparameter settings in the given `optimizer` (#83159 )
  * refactor `SequentialLR` & `ChainedScheduler` to support `ReduceLROnPlateau` with arbitrary step kwargs inputs.
    * #110761 
    * #68978
* Deprecate `step(epoch)` warnings and remove supports of `epoch`, and *probably change to closed-form implementation* as well.
  * conflicts: 
    * `Secheduelr` composed by a `SequentialLR`, will raise #119168. Closed-form implementation will cause every .
    * Currently, `CyclicLR` is missing closed-form implementation, as well as `ReduceLROnPlateau`.
    * Unittests for `SequentialLR` expect all schedulers to use close-form implementation, whereas the `ChainedScheduler` is not.
  * #55135
  * #67586
  * #116776
  * #67958
* Allow `param_groups` to be passed to `lr_scheduler` class
  * #72146
  * probable issues:
    * `param_groups` must be an instance of `dict` right from the beginning, or else any additional keys added by the `optimizer` will not be accessible by the `lr_scheduler` class.
    * Currently, there's no way to validate whether the `param_groups` passed to the `lr_scheduler` class are all in the same `optimizer` instance or not.
* Deprecate `print_lr`
  * #99270


## Proposed Changes
1. A new abstract class `_SchedulerBase` will be created.
```python
class _SchedulerBase(abc.ABC):
    optimizer: Optimizer

    @abc.abstractmethod
    def step(self, **kwargs):
        raise NotImplementedError

    @abc.abstractmethod
    def state_dict(self): ...

    @abc.abstractmethod
    def load_state_dict(self): ...
```
2. Reneme `lr_scheduler` to `Scheduler` and refactor it to support all hyperparameter settings in the given `optimizer`.
```python
class Scheduler(_SchedulerBase):
    @property
    def targets(self) -> Sequence[str]:
        """The targets that the schedulers can update."""
        ...

    @property
    def last_targets(self) -> Sequence[Dict[str, Any]]:
        """The targets at the last step for each parameter group."""
        ...

    def update_targets(self, **kwargs) -> None: 
        """Update the targets (lr, momentum, ...) for each parameter group."""
        ...
    
    def step(self, epoch: Optional[int] = None, **kwargs) -> None:
        """Update the learning rate of each parameter group."""
        if epoch is not None:
            warnings.warn("The epoch parameter is deprecated and will be removed in the future.", DeprecationWarning)
            self.last_step = epoch
        else:
          self.last_step += 1
        
        with _enable_get_lr_call(self):
            self.update_targets(step=self.last_step, **kwargs)
```
3. Refactor `SequentialLR` & `ChainedScheduler` to support `ReduceLROnPlateau` with arbitrary step kwargs inputs.
```python
class ComposedScheduler(_SchedulerBase, abc.ABC):
    """ChainedScheduler and SequentialLR can subclass this."""

    def __init__(
            self,
            schedulers: Sequence[LRScheduler],
            last_step: int = -1,
    ): ...

    def step(self, epoch: Optional[int] = None, **kwargs):
      # check epoch deprecation
      self.step_schedulers(step=self.last_step, **kwargs)

    @abc.abstractmethod
    def step_schedulers(self, **kwargs): ...

    @property
    def optimizer(self) -> Optimizer:
      return self.schedulers[0].optimizer

```

I like to spend some time fixing some of the issues regarding to the current scheduler module, and feel free to discuss with me further.

### Alternatives

leave it as it is.

### Additional context

Current development status: [link](https://github.com/GdoongMathew/pytorch/blob/refactor/lr_sch_interface/torch/optim/lr_scheduler.py)

cc @vincentqb @jbschlosser @albanD @janeyx99 @crcrpar