Fix TypeError in configure_optimizers when running with ReduceLROnPlateau

examples/fabric/tensor_parallel/train.py

@@ -1,14 +1,13 @@
 import lightning as L
 import torch
 import torch.nn.functional as F
+from data import RandomTokenDataset
 from lightning.fabric.strategies import ModelParallelStrategy
 from model import ModelArgs, Transformer
 from parallelism import parallelize
 from torch.distributed.tensor.parallel import loss_parallel
 from torch.utils.data import DataLoader
 
-from data import RandomTokenDataset
-
 
 def train():
     strategy = ModelParallelStrategy(

examples/pytorch/tensor_parallel/train.py

@@ -1,14 +1,13 @@
 import lightning as L
 import torch
 import torch.nn.functional as F
+from data import RandomTokenDataset
 from lightning.pytorch.strategies import ModelParallelStrategy
 from model import ModelArgs, Transformer
 from parallelism import parallelize
 from torch.distributed.tensor.parallel import loss_parallel
 from torch.utils.data import DataLoader
 
-from data import RandomTokenDataset
-
 
 class Llama3(L.LightningModule):
     def __init__(self):

src/lightning/pytorch/utilities/types.py

@@ -104,18 +104,25 @@ class LRSchedulerConfigType(TypedDict, total=False):
     strict: bool
 
 
-class OptimizerLRSchedulerConfig(TypedDict):
+class OptimizerConfigDict(TypedDict):
     optimizer: Optimizer
-    lr_scheduler: NotRequired[Union[LRSchedulerTypeUnion, LRSchedulerConfigType]]
+
+
+class OptimizerLRSchedulerConfigDict(TypedDict):
+    optimizer: Optimizer
+    lr_scheduler: Union[LRSchedulerTypeUnion, LRSchedulerConfigType]
+    monitor: NotRequired[str]
 
 
 OptimizerLRScheduler = Optional[
     Union[
         Optimizer,
         Sequence[Optimizer],
         tuple[Sequence[Optimizer], Sequence[Union[LRSchedulerTypeUnion, LRSchedulerConfig]]],
-        OptimizerLRSchedulerConfig,
-        Sequence[OptimizerLRSchedulerConfig],
+        OptimizerConfigDict,
+        OptimizerLRSchedulerConfigDict,
+        Sequence[OptimizerConfigDict],
+        Sequence[OptimizerLRSchedulerConfigDict],
     ]
 ]