Cannot import OptimizerLRSchedulerConfig or OptimizerLRSchedulerConfigDict

[zordi-youngsun]

Bug description

Since I bumped up lightning to 2.5.0, the configure_optimizers has been failing the type checker. I saw that OptimizerLRSchedulerConfig had been replaced with OptimizerLRSchedulerConfigDict, but I cannot import any of them.

What version are you seeing the problem on?

v2.5

How to reproduce the bug

import torch
import pytorch_lightning as pl
from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfigDict
from torch.optim.lr_scheduler import ReduceLROnPlateau

class Model(pl.LightningModule):
    ...
    def configure_optimizers(self) -> OptimizerLRSchedulerConfigDict:
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        scheduler = ReduceLROnPlateau(
            optimizer, mode="min", factor=0.1, patience=20, min_lr=1e-6
        )
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "monitor": "val_loss",
                "interval": "epoch",
                "frequency": 1,
            },
        }

Error messages and logs

In [2]: import lightning

In [3]: lightning.__version__
Out[3]: '2.5.0'

In [4]: from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfigDict
---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
Cell In[4], line 1
----> 1 from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfigDict

ImportError: cannot import name 'OptimizerLRSchedulerConfigDict' from 'lightning.pytorch.utilities.types' (/home/test/.venv/lib/python3.11/site-packages/lightning/pytorch/utilities/types.py)

In [5]: from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfig
---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
Cell In[5], line 1
----> 1 from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfig

ImportError: cannot import name 'OptimizerLRSchedulerConfig' from 'lightning.pytorch.utilities.types' (/home/test/.venv/lib/python3.11/site-packages/lightning/pytorch/utilities/types.py)

Environment

Current environment

#- PyTorch Lightning Version (e.g., 2.5.0):
#- PyTorch Version (e.g., 2.5):
#- Python version (e.g., 3.12):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):

More info

No response

[lantiga]

Thanks for the heads up, weird we didn't catch it earlier. We'll issue a post asap.

[lantiga]

Ok, two things:

• (less important): in your example you are importing the pytorch_lightning package, but the other type from lightning.pytorch. You should pick one or the other or you might get surprises.

• the return type for configure_optimizers hasn't been actually changed, it was OptimizerLRScheduler before and it still is. However that is a composite type and I renamed it in a PR: it was OptimizerLRSchedulerConfig and it now is OptimizerLRSchedulerConfigType thinking that it would not be impactful downstream but apparently it is.

So you can fix the above by doing

from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfigType

but I'm thinking we could just quickly go back to the previous name and avoid breaking the contract

[zordi-youngsun]

Thanks for your quick response. I replaced all instances of lightning.pytorch with pytorch_lightning (since I could only access the fit method through pytorch_lightning on my end) and confirmed that the config type you suggested works.

My model is now running successfully with the following import:

from pytorch_lightning.utilities.types import OptimizerLRSchedulerConfigType

If you are going back to the previous name, I’ll wait for the next release.

[lantiga]

Release is out: https://github.com/Lightning-AI/pytorch-lightning/releases/tag/2.5.0.post0

Previous name was restored