Add log_key_prefix parameter to LearningRateMonitor (#21612)

* Add `log_key_prefix` parameter to `LearningRateMonitor` callback

Allow users to prepend a configurable prefix to all metric names logged
by `LearningRateMonitor`. This is useful for grouping learning rate
metrics in loggers like TensorBoard (e.g., `optim/lr-Adam` instead of
`lr-Adam`).

Fixes #21590

* Update src/lightning/pytorch/CHANGELOG.md

---------

Co-authored-by: Deependu <deependujha21@gmail.com>

Author: Krishnachaitanyakc

src/lightning/pytorch/CHANGELOG.md

@@ -10,7 +10,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
--
+- Added `log_key_prefix` parameter to `LearningRateMonitor` callback for prefixing logged metric names ([#21612](https://github.com/Lightning-AI/pytorch-lightning/issues/21612))
 
 ### Changed
 

src/lightning/pytorch/callbacks/lr_monitor.py

@@ -46,6 +46,10 @@ class LearningRateMonitor(Callback):
             has the ``momentum`` or ``betas`` attribute. Defaults to ``False``.
         log_weight_decay: option to also log the weight decay values of the optimizer. Defaults to
             ``False``.
+        log_key_prefix: optional string prefix to prepend to all logged metric names. Useful for
+            grouping learning rate metrics in loggers like TensorBoard. For example, setting
+            ``log_key_prefix="optim/"`` would log ``optim/lr-Adam`` instead of ``lr-Adam``.
+            Defaults to ``None`` (no prefix).
 
     Raises:
         MisconfigurationException:
@@ -96,13 +100,15 @@ def __init__(
         logging_interval: Optional[Literal["step", "epoch"]] = None,
         log_momentum: bool = False,
         log_weight_decay: bool = False,
+        log_key_prefix: Optional[str] = None,
     ) -> 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.log_key_prefix = log_key_prefix or ""
 
         self.lrs: dict[str, list[float]] = {}
         self.last_momentum_values: dict[str, Optional[list[float]]] = {}
@@ -361,4 +367,7 @@ def _check_duplicates_and_update_name(
             )
 
         name = self._add_prefix(name, optimizer_cls, seen_optimizer_types)
-        return [self._add_suffix(name, param_groups, i) for i in range(len(param_groups))]
+        names = [self._add_suffix(name, param_groups, i) for i in range(len(param_groups))]
+        if self.log_key_prefix:
+            names = [f"{self.log_key_prefix}{n}" for n in names]
+        return names

tests/tests_pytorch/callbacks/test_lr_monitor.py

@@ -709,3 +709,111 @@ def configure_optimizers(self):
     # Verify the callback metric tensor was created successfully
     assert "lr-SGD" in trainer.callback_metrics
     assert isinstance(trainer.callback_metrics["lr-SGD"], torch.Tensor)
+
+
+def test_lr_monitor_log_key_prefix(tmp_path):
+    """Test that learning rate metric names are correctly prefixed when log_key_prefix is set."""
+    model = BoringModel()
+
+    lr_monitor = LearningRateMonitor(log_key_prefix="optim/")
+    trainer = Trainer(
+        default_root_dir=tmp_path,
+        max_epochs=2,
+        limit_val_batches=0.1,
+        limit_train_batches=0.5,
+        callbacks=[lr_monitor],
+        logger=CSVLogger(tmp_path),
+    )
+    trainer.fit(model)
+
+    assert lr_monitor.lrs, "No learning rates logged"
+    assert list(lr_monitor.lrs) == ["optim/lr-SGD"]
+    assert "optim/lr-SGD" in trainer.callback_metrics
+
+
+def test_lr_monitor_log_key_prefix_with_momentum_and_weight_decay(tmp_path):
+    """Test that prefix is applied to momentum and weight decay metric names as well."""
+
+    class CustomModel(BoringModel):
+        def configure_optimizers(self):
+            optimizer = optim.Adam(self.parameters(), lr=1e-2, betas=(0.9, 0.999), weight_decay=0.01)
+            lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1)
+            return [optimizer], [lr_scheduler]
+
+    model = CustomModel()
+    lr_monitor = LearningRateMonitor(log_momentum=True, log_weight_decay=True, log_key_prefix="train/")
+    trainer = Trainer(
+        default_root_dir=tmp_path,
+        max_epochs=2,
+        limit_val_batches=2,
+        limit_train_batches=5,
+        log_every_n_steps=1,
+        callbacks=[lr_monitor],
+        logger=CSVLogger(tmp_path),
+    )
+    trainer.fit(model)
+
+    assert list(lr_monitor.lrs) == ["train/lr-Adam"]
+    assert all(k == "train/lr-Adam-momentum" for k in lr_monitor.last_momentum_values)
+    assert all(k == "train/lr-Adam-weight_decay" for k in lr_monitor.last_weight_decay_values)
+
+
+def test_lr_monitor_log_key_prefix_multi_optimizers(tmp_path):
+    """Test that prefix is applied correctly with multiple optimizers."""
+
+    class MultiOptModel(BoringModel):
+        def __init__(self):
+            super().__init__()
+            self.automatic_optimization = False
+
+        def training_step(self, batch, batch_idx):
+            opt1, opt2 = self.optimizers()
+
+            loss = self.loss(self.step(batch))
+            opt1.zero_grad()
+            self.manual_backward(loss)
+            opt1.step()
+
+            loss = self.loss(self.step(batch))
+            opt2.zero_grad()
+            self.manual_backward(loss)
+            opt2.step()
+
+        def configure_optimizers(self):
+            optimizer1 = optim.Adam(self.parameters(), lr=1e-2)
+            optimizer2 = optim.SGD(self.parameters(), lr=1e-2)
+            return [optimizer1, optimizer2]
+
+    model = MultiOptModel()
+    lr_monitor = LearningRateMonitor(log_key_prefix="hparams/")
+    trainer = Trainer(
+        default_root_dir=tmp_path,
+        max_epochs=2,
+        limit_val_batches=0.1,
+        limit_train_batches=5,
+        log_every_n_steps=1,
+        callbacks=[lr_monitor],
+        logger=CSVLogger(tmp_path),
+    )
+    trainer.fit(model)
+
+    assert lr_monitor.lrs, "No learning rates logged"
+    assert list(lr_monitor.lrs) == ["hparams/lr-Adam", "hparams/lr-SGD"]
+
+
+def test_lr_monitor_log_key_prefix_none(tmp_path):
+    """Test that when log_key_prefix is None (default), metric names are unchanged."""
+    model = BoringModel()
+
+    lr_monitor = LearningRateMonitor(log_key_prefix=None)
+    trainer = Trainer(
+        default_root_dir=tmp_path,
+        max_epochs=2,
+        limit_val_batches=0.1,
+        limit_train_batches=0.5,
+        callbacks=[lr_monitor],
+        logger=CSVLogger(tmp_path),
+    )
+    trainer.fit(model)
+
+    assert list(lr_monitor.lrs) == ["lr-SGD"]