Fix LearningRateMonitor crash for optimizers without weight decay

src/lightning/pytorch/CHANGELOG.md

@@ -36,6 +36,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed `LearningRateMonitor` raising `KeyError` with `log_weight_decay=True` for optimizers without weight decay, such as `Rprop` ([#22035](https://github.com/Lightning-AI/pytorch-lightning/issues/22035))
+
 - Fixed `LightningCLI` emitting `jsonargparse` deprecation warnings ([#21900](https://github.com/Lightning-AI/pytorch-lightning/issues/21900))
 
 - Fixed `RichProgressBar` showing a nonsensical negative epoch total (e.g. `Epoch 5/-2`) when `Trainer(max_epochs=-1)` (unlimited epochs) is used and training stops via another condition ([#21925](https://github.com/Lightning-AI/pytorch-lightning/issues/21925))

src/lightning/pytorch/callbacks/lr_monitor.py

@@ -129,22 +129,25 @@ def on_train_start(self, trainer: "pl.Trainer", *args: Any, **kwargs: Any) -> No
                 "Cannot use `LearningRateMonitor` callback with `Trainer` that has no logger."
             )
 
-        if self.log_momentum:
+        def _check_no_key(key: str) -> bool:
+            if trainer.lr_scheduler_configs:
+                return any(key not in config.scheduler.optimizer.defaults for config in trainer.lr_scheduler_configs)
 
-            def _check_no_key(key: str) -> bool:
-                if trainer.lr_scheduler_configs:
-                    return any(
-                        key not in config.scheduler.optimizer.defaults for config in trainer.lr_scheduler_configs
-                    )
+            return any(key not in optimizer.defaults for optimizer in trainer.optimizers)
 
-                return any(key not in optimizer.defaults for optimizer in trainer.optimizers)
+        if self.log_momentum and _check_no_key("momentum") and _check_no_key("betas"):
+            rank_zero_warn(
+                "You have set log_momentum=True, but some optimizers do not"
+                " have momentum. This will log a value 0 for the momentum.",
+                category=RuntimeWarning,
+            )
 
-            if _check_no_key("momentum") and _check_no_key("betas"):
-                rank_zero_warn(
-                    "You have set log_momentum=True, but some optimizers do not"
-                    " have momentum. This will log a value 0 for the momentum.",
-                    category=RuntimeWarning,
-                )
+        if self.log_weight_decay and _check_no_key("weight_decay"):
+            rank_zero_warn(
+                "You have set log_weight_decay=True, but some optimizers do not"
+                " have weight decay. Weight decay will not be logged for them.",
+                category=RuntimeWarning,
+            )
 
         # Find names for schedulers
         names: list[list[str]] = []
@@ -267,7 +270,7 @@ def _extract_momentum(self, param_group: dict[str, list[float]], name: str, use_
 
     def _extract_weight_decay(self, param_group: dict[str, Any], name: str) -> dict[str, Any]:
         """Extracts the weight decay statistics from a parameter group."""
-        if not self.log_weight_decay:
+        if not self.log_weight_decay or "weight_decay" not in param_group:
             return {}
 
         weight_decay = param_group["weight_decay"]

tests/tests_pytorch/callbacks/test_lr_monitor.py

@@ -122,6 +122,74 @@ def configure_optimizers(self):
     assert all(k == "lr-ASGD-momentum" for k in lr_monitor.last_momentum_values)
 
 
+def test_log_weight_decay_no_weight_decay_optimizer(tmp_path):
+    """Test that weight decay is not logged for an optimizer without it, and a warning is raised."""
+
+    class LogWeightDecayModel(BoringModel):
+        def configure_optimizers(self):
+            optimizer = optim.Rprop(self.parameters(), lr=1e-2)
+            lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1)
+            return [optimizer], [lr_scheduler]
+
+    model = LogWeightDecayModel()
+    lr_monitor = LearningRateMonitor(log_weight_decay=True)
+    trainer = Trainer(
+        default_root_dir=tmp_path,
+        max_epochs=1,
+        limit_val_batches=2,
+        limit_train_batches=5,
+        log_every_n_steps=1,
+        callbacks=[lr_monitor],
+        logger=CSVLogger(tmp_path),
+    )
+    with pytest.warns(RuntimeWarning, match="optimizers do not have weight decay."):
+        trainer.fit(model)
+
+    assert lr_monitor.last_weight_decay_values == {"lr-Rprop-weight_decay": None}
+    assert "lr-Rprop-weight_decay" not in trainer.callback_metrics
+
+
+def test_log_weight_decay_mixed_optimizers(tmp_path):
+    """Test that weight decay is still logged for the optimizer that has it when another optimizer does not."""
+
+    class LogWeightDecayModel(BoringModel):
+        def __init__(self):
+            super().__init__()
+            self.automatic_optimization = False
+
+        def training_step(self, batch, batch_idx):
+            for opt in self.optimizers():
+                loss = self.loss(self.step(batch))
+                opt.zero_grad()
+                self.manual_backward(loss)
+                opt.step()
+
+        def configure_optimizers(self):
+            optimizer1 = optim.AdamW(self.parameters(), lr=1e-2, weight_decay=0.1)
+            optimizer2 = optim.Rprop(self.parameters(), lr=1e-2)
+            lr_scheduler1 = optim.lr_scheduler.StepLR(optimizer1, step_size=1)
+            lr_scheduler2 = optim.lr_scheduler.StepLR(optimizer2, step_size=1)
+            return [optimizer1, optimizer2], [lr_scheduler1, lr_scheduler2]
+
+    model = LogWeightDecayModel()
+    lr_monitor = LearningRateMonitor(log_weight_decay=True)
+    trainer = Trainer(
+        default_root_dir=tmp_path,
+        max_epochs=1,
+        limit_val_batches=2,
+        limit_train_batches=5,
+        log_every_n_steps=1,
+        callbacks=[lr_monitor],
+        logger=CSVLogger(tmp_path),
+    )
+    with pytest.warns(RuntimeWarning, match="optimizers do not have weight decay."):
+        trainer.fit(model)
+
+    assert lr_monitor.last_weight_decay_values == {"lr-AdamW-weight_decay": 0.1, "lr-Rprop-weight_decay": None}
+    assert trainer.callback_metrics["lr-AdamW-weight_decay"] == 0.1
+    assert "lr-Rprop-weight_decay" not in trainer.callback_metrics
+
+
 def test_lr_monitor_no_lr_scheduler_single_lr(tmp_path):
     """Test that learning rates are extracted and logged for no lr scheduler."""