Repository: Lightning-AI/pytorch-lightning
Pull request: Add `log_weight_decay` argument to `LearningRateMonitor` callback
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

Fixes #18117

<!-- Does your PR introduce any breaking changes? If yes, please list them. -->

<details>
  <summary><b>Before submitting</b></summary>

- Was this **discussed/agreed** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- Did you make sure to **update the documentation** with your changes? (if necessary)
- Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- Did you list all the **breaking changes** introduced by this pull request?
- Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

</details>

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

<details>
  <summary>Reviewer checklist</summary>

- [ ] Is this pull request ready for review? (if not, please submit in draft mode)
- [ ] Check that all items from **Before submitting** are resolved
- [ ] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

</details>

<!--

Did you have fun?

Make sure you had fun coding 🙃

-->


Problem statement:
Add `log_weight_decay` option to `LearningRateMonitor` callback
### Description & Motivation

In more sophisticated training settings, we usually apply a dedicated weight decay scheduler (for example DINO). Thus, it will be great to monitor these statistics, optionally, from the `LearningRateMonitor` callback.

### Pitch

A solution should look something like the following, but build in the LearningRateMonitor:
```py
"""Learning rate monitor logger callback."""
from typing import Any, Dict, List

from pytorch_lightning import callbacks
from torch import optim


class LearningRateMonitor(callbacks.LearningRateMonitor):
    """Extends the PL version with the ability to log the weight decay."""

    def __init__(
        self,
        logging_interval: str | None = None,
        log_momentum: bool = False,
        log_weight_decay: bool = False,
    ) -> None:
        """Initializes the callback."""
        super().__init__(logging_interval, log_momentum)
        self.log_weight_decay = log_weight_decay

    def _get_lr_momentum_stat(
        self, optimizer: optim.Optimizer, names: List[str]
    ) -> Dict[str, float]:
        """Gathers and returns the learning rate, momentum and weight decay statistics."""
        lr_momentum_wd_stat = {}
        param_groups = optimizer.param_groups
        use_betas = "betas" in optimizer.defaults

        for pg, name in zip(param_groups, names):
            lr = self._extract_lr(pg, name)
            lr_momentum_wd_stat.update(lr)
            momentum = self._extract_momentum(
                param_group=pg, name=name.replace(name, f"{name}-momentum"), use_betas=use_betas
            )
            lr_momentum_wd_stat.update(momentum)
            weight_decay = self._extract_weight_decay(pg, name + "-weight_decay")
            lr_momentum_wd_stat.update(weight_decay)

        return lr_momentum_wd_stat

    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:
            return {}

        wd = param_group["weight_decay"]
        return {name: wd}
```

### Alternatives

_No response_

### Additional context

_No response_

cc @borda

Hints from issue discussion:
Thanks @lobantseff for this suggestion. I don't see what could stand in the way of such an addition. A PR for this would be welcomed if you are interested to contribute the feature :)

Commit message: Merge branch 'master' into ioannis@add-log_weight_decay-option-to-LearningRateMonitor-callback

Patch to review:
diff --git a/src/lightning/pytorch/callbacks/lr_monitor.py b/src/lightning/pytorch/callbacks/lr_monitor.py
index d823cf5d52010..1aca99e718554 100644
--- a/src/lightning/pytorch/callbacks/lr_monitor.py
+++ b/src/lightning/pytorch/callbacks/lr_monitor.py
@@ -21,7 +21,7 @@
 """
 import itertools
 from collections import defaultdict
-from typing import Any, DefaultDict, Dict, List, Optional, Set, Tuple, Type
+from typing import Any, DefaultDict, Dict, List, Literal, Optional, Set, Tuple, Type
 
 import torch
 from torch.optim.optimizer import Optimizer
@@ -87,13 +87,22 @@ def configure_optimizer(self):
 
     """
 
-    def __init__(self, logging_interval: Optional[str] = None, log_momentum: bool = False) -> None:
+    def __init__(
+        self,
+        logging_interval: Optional[Literal["step", "epoch"]] = None,
+        log_momentum: bool = False,
+        log_weight_decay: bool = False,
+    ) -> None:
         if logging_interval not in (None, "step", "epoch"):
             raise MisconfigurationException("logging_interval should be `step` or `epoch` or `None`.")
 
         self.logging_interval = logging_interval
         self.log_momentum = log_momentum
+        self.log_weight_decay = log_weight_decay
+
         self.lrs: Dict[str, List[float]] = {}
+        self.last_momentum_values: Dict[str, List[float]] = {}
+        self.last_weight_decay_values: Dict[str, List[float]] = {}
 
     def on_train_start(self, trainer: "pl.Trainer", *args: Any, **kwargs: Any) -> None:
         """Called before training, determines unique names for all lr schedulers in the case of multiple of the same
@@ -147,6 +156,7 @@ def _check_no_key(key: str) -> bool:
         names_flatten = list(itertools.chain.from_iterable(names))
         self.lrs = {name: [] for name in names_flatten}
         self.last_momentum_values = {name + "-momentum": None for name in names_flatten}
+        self.last_weight_decay_values = {name + "-weight_decay": None for name in names_flatten}
 
     def on_train_batch_start(self, trainer: "pl.Trainer", *args: Any, **kwargs: Any) -> None:
         if not trainer._logger_connector.should_update_logs:
@@ -182,7 +192,7 @@ def _extract_stats(self, trainer: "pl.Trainer", interval: str) -> Dict[str, floa
         for name, config in zip(scheduler_hparam_keys, trainer.lr_scheduler_configs):
             if interval in [config.interval, "any"]:
                 opt = config.scheduler.optimizer
-                current_stat = self._get_lr_momentum_stat(opt, name)
+                current_stat = self._get_optimizer_stats(opt, name)
                 latest_stat.update(current_stat)
 
         optimizer_hparam_keys, optimizers_without_scheduler = self._find_names_from_optimizers(
@@ -193,7 +203,7 @@ def _extract_stats(self, trainer: "pl.Trainer", interval: str) -> Dict[str, floa
         self._remap_keys(optimizer_hparam_keys)
 
         for opt, names in zip(optimizers_without_scheduler, optimizer_hparam_keys):
-            current_stat = self._get_lr_momentum_stat(opt, names)
+            current_stat = self._get_optimizer_stats(opt, names)
             latest_stat.update(current_stat)
 
         trainer.callback_metrics.update(
@@ -202,20 +212,22 @@ def _extract_stats(self, trainer: "pl.Trainer", interval: str) -> Dict[str, floa
 
         return latest_stat
 
-    def _get_lr_momentum_stat(self, optimizer: Optimizer, names: List[str]) -> Dict[str, float]:
-        lr_momentum_stat = {}
+    def _get_optimizer_stats(self, optimizer: Optimizer, names: List[str]) -> Dict[str, float]:
+        stats = {}
         param_groups = optimizer.param_groups
         use_betas = "betas" in optimizer.defaults
 
         for pg, name in zip(param_groups, names):
             lr = self._extract_lr(pg, name)
-            lr_momentum_stat.update(lr)
+            stats.update(lr)
             momentum = self._extract_momentum(
                 param_group=pg, name=name.replace(name, f"{name}-momentum"), use_betas=use_betas
             )
-            lr_momentum_stat.update(momentum)
+            stats.update(momentum)
+            weight_decay = self._extract_weight_decay(pg, name + "-weight_decay")
+            stats.update(weight_decay)
 
-        return lr_momentum_stat
+        return stats
 
     def _extract_lr(self, param_group: Dict[str, Any], name: str) -> Dict[str, Any]:
         lr = param_group["lr"]
@@ -240,6 +252,15 @@ def _extract_momentum(self, param_group: Dict[str, List], name: str, use_betas:
         self.last_momentum_values[name] = momentum
         return {name: momentum}
 
+    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:
+            return {}
+
+        weight_decay = param_group["weight_decay"]
+        self.last_weight_decay_values[name] = weight_decay
+        return {name: weight_decay}
+
     def _add_prefix(
         self, name: str, optimizer_cls: Type[Optimizer], seen_optimizer_types: DefaultDict[Type[Optimizer], int]
     ) -> str:
diff --git a/tests/tests_pytorch/callbacks/test_lr_monitor.py b/tests/tests_pytorch/callbacks/test_lr_monitor.py
index 7f89ea845b819..9a5ba3b44b54a 100644
--- a/tests/tests_pytorch/callbacks/test_lr_monitor.py
+++ b/tests/tests_pytorch/callbacks/test_lr_monitor.py
@@ -50,7 +50,7 @@ def test_lr_monitor_single_lr(tmpdir):
 
 @pytest.mark.parametrize("opt", ["SGD", "Adam"])
 def test_lr_monitor_single_lr_with_momentum(tmpdir, opt: str):
-    """Test that learning rates and momentum are extracted and logged for single lr scheduler."""
+    """Test that learning rates, momentum and weight decay are extracted and logged for single lr scheduler."""
 
     class LogMomentumModel(BoringModel):
         def __init__(self, opt):
@@ -68,7 +68,7 @@ def configure_optimizers(self):
             return [optimizer], [lr_scheduler]
 
     model = LogMomentumModel(opt=opt)
-    lr_monitor = LearningRateMonitor(log_momentum=True)
+    lr_monitor = LearningRateMonitor(log_momentum=True, log_weight_decay=True)
     trainer = Trainer(
         default_root_dir=tmpdir,
         max_epochs=2,
@@ -84,6 +84,12 @@ def configure_optimizers(self):
     assert len(lr_monitor.last_momentum_values) == len(trainer.lr_scheduler_configs)
     assert all(k == f"lr-{opt}-momentum" for k in lr_monitor.last_momentum_values)
 
+    assert all(
+        v is not None for v in lr_monitor.last_weight_decay_values.values()
+    ), "Expected weight decay to be logged"
+    assert len(lr_monitor.last_weight_decay_values) == len(trainer.lr_scheduler_configs)
+    assert all(k == f"lr-{opt}-weight_decay" for k in lr_monitor.last_weight_decay_values)
+
 
 def test_log_momentum_no_momentum_optimizer(tmpdir):
     """Test that if optimizer doesn't have momentum then a warning is raised with log_momentum=True."""
@@ -599,12 +605,18 @@ def forward(self, x):
 
         def configure_optimizers(self):
             param_groups = [
-                {"params": list(self.linear_a.parameters())},
-                {"params": list(self.linear_b.parameters())},
+                {
+                    "params": list(self.linear_a.parameters()),
+                    "weight_decay": 0.1,
+                },
+                {
+                    "params": list(self.linear_b.parameters()),
+                    "weight_decay": 0.1,
+                },
             ]
             return torch.optim.Adam(param_groups, lr=self.hparams.lr, betas=self.hparams.momentum)
 
-    lr_monitor = LearningRateMonitor(log_momentum=True)
+    lr_monitor = LearningRateMonitor(log_momentum=True, log_weight_decay=True)
     trainer = Trainer(
         default_root_dir=tmpdir,
         max_epochs=2,
@@ -618,6 +630,7 @@ def configure_optimizers(self):
 
     lr = 1e-2
     momentum = 0.7
+    weight_decay = 0.1
     model = TestModel(lr=lr, momentum=(momentum, 0.999))
     trainer.fit(model)
 
@@ -625,6 +638,8 @@ def configure_optimizers(self):
     assert list(lr_monitor.lrs) == ["lr-Adam/pg1", "lr-Adam/pg2"]
     assert list(lr_monitor.last_momentum_values) == ["lr-Adam/pg1-momentum", "lr-Adam/pg2-momentum"]
     assert all(val == momentum for val in lr_monitor.last_momentum_values.values())
+    assert list(lr_monitor.last_weight_decay_values) == ["lr-Adam/pg1-weight_decay", "lr-Adam/pg2-weight_decay"]
+    assert all(val == weight_decay for val in lr_monitor.last_weight_decay_values.values())
     assert all(all(val == lr for val in lr_monitor.lrs[lr_key]) for lr_key in lr_monitor.lrs)
 
 


Code review comments:
diff_hunk: @@ -202,20 +212,22 @@ def _extract_stats(self, trainer: "pl.Trainer", interval: str) -> Dict[str, floa
 
         return latest_stat
 
-    def _get_lr_momentum_stat(self, optimizer: Optimizer, names: List[str]) -> Dict[str, float]:
-        lr_momentum_stat = {}
+    def _get_optimizer_stats(self, optimizer: Optimizer, names: List[str]) -> Dict[str, float]:
+        stats = {}
         param_groups = optimizer.param_groups
         use_betas = "betas" in optimizer.defaults
 
         for pg, name in zip(param_groups, names):
             lr = self._extract_lr(pg, name)
-            lr_momentum_stat.update(lr)
+            stats.update(lr)
             momentum = self._extract_momentum(
                 param_group=pg, name=name.replace(name, f"{name}-momentum"), use_betas=use_betas
             )
-            lr_momentum_stat.update(momentum)
+            stats.update(momentum)
+            weight_decay = self._extract_weight_decay(pg, name + "-weight_decay")
path: src/lightning/pytorch/callbacks/lr_monitor.py
text: @user1:
nit: we use fstring directly above

```suggestion
            weight_decay = self._extract_weight_decay(pg, f"{name}-weight_decay")
```