From 24f0e31dd847f6d5faadc11ab4f40b6aee804c0a Mon Sep 17 00:00:00 2001 From: Raphael Fakhri <153192858+RaphaelFakhri@users.noreply.github.com> Date: Tue, 20 Oct 2026 09:43:20 +0000 Subject: [PATCH 1/2] Fix LearningRateMonitor crash for optimizers without weight decay --- src/lightning/pytorch/callbacks/lr_monitor.py | 2 +- .../callbacks/test_lr_monitor.py | 26 +++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/src/lightning/pytorch/callbacks/lr_monitor.py b/src/lightning/pytorch/callbacks/lr_monitor.py index e830374840f24..f60c792448b49 100644 --- a/src/lightning/pytorch/callbacks/lr_monitor.py +++ b/src/lightning/pytorch/callbacks/lr_monitor.py @@ -270,7 +270,7 @@ def _extract_weight_decay(self, param_group: dict[str, Any], name: str) -> dict[ if not self.log_weight_decay: return {} - weight_decay = param_group["weight_decay"] + weight_decay = param_group.get("weight_decay", 0) self.last_weight_decay_values[name] = weight_decay return {name: weight_decay} diff --git a/tests/tests_pytorch/callbacks/test_lr_monitor.py b/tests/tests_pytorch/callbacks/test_lr_monitor.py index 01bddd8234f8c..a7131f66d7d3c 100644 --- a/tests/tests_pytorch/callbacks/test_lr_monitor.py +++ b/tests/tests_pytorch/callbacks/test_lr_monitor.py @@ -122,6 +122,32 @@ 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 an optimizer without weight decay logs a value of 0 with log_weight_decay=True.""" + + 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), + ) + trainer.fit(model) + + assert all(v == 0 for v in lr_monitor.last_weight_decay_values.values()), "Expected weight decay to be logged" + assert list(lr_monitor.last_weight_decay_values) == ["lr-Rprop-weight_decay"] + + def test_lr_monitor_no_lr_scheduler_single_lr(tmp_path): """Test that learning rates are extracted and logged for no lr scheduler.""" From 95bbc7e9f85c3bf0388e0acaf64fc92aaa924bea Mon Sep 17 00:00:00 2001 From: Raphael Fakhri <153192858+RaphaelFakhri@users.noreply.github.com> Date: Fri, 30 Oct 2026 15:02:13 +0000 Subject: [PATCH 2/2] Skip weight decay for param groups without it and warn once --- src/lightning/pytorch/CHANGELOG.md | 2 + src/lightning/pytorch/callbacks/lr_monitor.py | 33 ++++++------ .../callbacks/test_lr_monitor.py | 50 +++++++++++++++++-- 3 files changed, 66 insertions(+), 19 deletions(-) diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md index 1c73ea347751a..90a65a144457e 100644 --- a/src/lightning/pytorch/CHANGELOG.md +++ b/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)) diff --git a/src/lightning/pytorch/callbacks/lr_monitor.py b/src/lightning/pytorch/callbacks/lr_monitor.py index f60c792448b49..a063347e4d046 100644 --- a/src/lightning/pytorch/callbacks/lr_monitor.py +++ b/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,10 +270,10 @@ 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.get("weight_decay", 0) + weight_decay = param_group["weight_decay"] self.last_weight_decay_values[name] = weight_decay return {name: weight_decay} diff --git a/tests/tests_pytorch/callbacks/test_lr_monitor.py b/tests/tests_pytorch/callbacks/test_lr_monitor.py index a7131f66d7d3c..e44d2c4942156 100644 --- a/tests/tests_pytorch/callbacks/test_lr_monitor.py +++ b/tests/tests_pytorch/callbacks/test_lr_monitor.py @@ -123,7 +123,7 @@ def configure_optimizers(self): def test_log_weight_decay_no_weight_decay_optimizer(tmp_path): - """Test that an optimizer without weight decay logs a value of 0 with log_weight_decay=True.""" + """Test that weight decay is not logged for an optimizer without it, and a warning is raised.""" class LogWeightDecayModel(BoringModel): def configure_optimizers(self): @@ -142,10 +142,52 @@ def configure_optimizers(self): callbacks=[lr_monitor], logger=CSVLogger(tmp_path), ) - trainer.fit(model) + 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 all(v == 0 for v in lr_monitor.last_weight_decay_values.values()), "Expected weight decay to be logged" - assert list(lr_monitor.last_weight_decay_values) == ["lr-Rprop-weight_decay"] + 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):