fix: log the convenience epoch metric once per step (#20902)

At an epoch boundary several metric groups are flushed at the same step
(the training-epoch-end flush, the validation-end flush, and the last
train-step flush). Each `step=None` flush re-stamped the convenience
`epoch` metric via `setdefault`, so loggers that keep per-call history
(e.g. MLflow) recorded two `epoch` datapoints per epoch and misaligned
epoch-based x-axes.

Track the last `(step, epoch)` for which the convenience `epoch` metric
was auto-added and skip re-adding it at the same step. The key is
`(step, epoch)` (not `step` alone) so an epoch that advances without an
optimizer step (e.g. under gradient accumulation) is still recorded. The
tracker is reset per run so a fresh fit/validate/test re-emits `epoch`.
Step-indexed loggers (TensorBoard, CSV) are unaffected because they merge
by step.

Update the logger-history assertions in test_all.py,
test_logger_connector.py, and the train/eval loop logging tests to expect
the single-`epoch`-per-step behavior, and add a focused regression test.

Author: gaurav0107

src/lightning/pytorch/CHANGELOG.md

@@ -24,6 +24,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed the convenience `epoch` metric being logged twice at the same step (once per metric group flushed at an epoch boundary), which produced duplicate `epoch` datapoints and misaligned epoch-based x-axes in loggers that keep per-call history such as MLflow ([#20902](https://github.com/Lightning-AI/pytorch-lightning/issues/20902))
+
 - Fixed PyTorch Lightning profiler not capturing dataloader worker initialization time ([#21771](https://github.com/Lightning-AI/pytorch-lightning/issues/21771))
 
 - Fixed `FSDPStrategy` raising `RuntimeError` under PyTorch 2.5+ when `root_device` is CPU, by passing an explicit `torch.device("cpu")` instead of `device_id=None` (relevant only when the GPU-accelerator guard is bypassed) ([#21774](https://github.com/Lightning-AI/pytorch-lightning/pull/21774))

src/lightning/pytorch/trainer/connectors/logger_connector/logger_connector.py

@@ -38,6 +38,9 @@ def __init__(self, trainer: "pl.Trainer") -> None:
         self._current_fx: Optional[str] = None
         # None: hasn't started, True: first loop iteration, False: subsequent iterations
         self._first_loop_iter: Optional[bool] = None
+        # tracks the last ``(step, epoch)`` for which the convenience ``epoch`` metric was
+        # auto-added, so it is not logged more than once at the same step (see ``log_metrics``)
+        self._logged_epoch_step: Optional[tuple[int, int]] = None
 
     def on_trainer_init(
         self,
@@ -121,9 +124,17 @@ def log_metrics(self, metrics: _OUT_DICT, step: Optional[int] = None) -> None:
             if step_metric is not None:
                 step = int(step_metric)
             else:
-                # added metrics for convenience
-                scalar_metrics.setdefault("epoch", self.trainer.current_epoch)
                 step = self.trainer.fit_loop.epoch_loop._batches_that_stepped
+                # add the convenience `epoch` metric, but only once per `(step, epoch)`.
+                # Several metric groups can be flushed at the same step (e.g. the
+                # training-epoch-end and validation-end flushes), and stamping `epoch` on
+                # each of them makes loggers that keep per-call history (e.g. MLflow) record
+                # duplicate `epoch` datapoints per epoch, misaligning epoch-based x-axes (#20902).
+                if "epoch" not in scalar_metrics:
+                    current_epoch = self.trainer.current_epoch
+                    if self._logged_epoch_step != (step, current_epoch):
+                        scalar_metrics["epoch"] = current_epoch
+                        self._logged_epoch_step = (step, current_epoch)
 
         # log actual metrics
         for logger in self.trainer.loggers:
@@ -230,6 +241,9 @@ def reset_metrics(self) -> None:
         self._progress_bar_metrics = {}
         self._logged_metrics = {}
         self._callback_metrics = {}
+        # reset per-run so the convenience `epoch` metric is re-emitted at the start of a
+        # fresh fit/validate/test run (this is called per-run, not per-epoch)
+        self._logged_epoch_step = None
 
     def reset_results(self) -> None:
         results = self.trainer._results

tests/tests_pytorch/loggers/test_all.py

@@ -125,14 +125,18 @@ def log_metrics(self, metrics, step):
     if logger_class == TensorBoardLogger:
         expected = [
             (0, ["epoch", "train_some_val"]),
-            (0, ["early_stop_on", "epoch", "val_loss"]),
+            # the convenience `epoch` metric is only stamped on the first flush at a
+            # given step, so the validation-end flush at step 0 no longer repeats it (#20902)
+            (0, ["early_stop_on", "val_loss"]),
             (1, ["epoch", "test_loss"]),
         ]
         assert log_metric_names == expected
     else:
         expected = [
             (0, ["epoch", "train_some_val"]),
-            (0, ["early_stop_on", "epoch", "val_loss"]),
+            # the convenience `epoch` metric is only stamped on the first flush at a
+            # given step, so the validation-end flush at step 0 no longer repeats it (#20902)
+            (0, ["early_stop_on", "val_loss"]),
             (1, ["epoch", "test_loss"]),
         ]
         assert log_metric_names == expected

tests/tests_pytorch/trainer/connectors/test_logger_connector.py

@@ -58,4 +58,5 @@ def test_uses_batches_that_stepped(mock_convert):
         metrics=mock_convert.return_value, step=trainer.fit_loop.epoch_loop._batches_that_stepped
     )
     logger.save.assert_called_once_with()
-    mock_convert.return_value.setdefault.assert_called_once_with("epoch", trainer.current_epoch)
+    # the convenience `epoch` metric is added once per `(step, epoch)` (see #20902)
+    mock_convert.return_value.__setitem__.assert_called_once_with("epoch", trainer.current_epoch)

tests/tests_pytorch/trainer/logging_/test_eval_loop_logging.py

@@ -543,16 +543,18 @@ def test_step(self, batch, batch_idx):
         "valid_loss_0",
         "valid_loss_1",
     }
+    # the validation-end flush shares a step with the earlier train flush, so it no longer
+    # re-emits a duplicate `epoch` (which the train flush already stamped) — see #20902.
     assert mock_log_metrics.mock_calls == [
         call({"hp_metric": -1}, None),
         call(metrics={"train_loss": ANY, "epoch": 0}, step=0),
         call(metrics={"valid_loss_0_step": ANY, "valid_loss_2": ANY}, step=0),
         call(metrics={"valid_loss_0_step": ANY, "valid_loss_2": ANY}, step=1),
-        call(metrics={"valid_loss_0_epoch": ANY, "valid_loss_1": ANY, "epoch": 0}, step=0),
+        call(metrics={"valid_loss_0_epoch": ANY, "valid_loss_1": ANY}, step=0),
         call(metrics={"train_loss": ANY, "epoch": 1}, step=1),
         call(metrics={"valid_loss_0_step": ANY, "valid_loss_2": ANY}, step=2),
         call(metrics={"valid_loss_0_step": ANY, "valid_loss_2": ANY}, step=3),
-        call(metrics={"valid_loss_0_epoch": ANY, "valid_loss_1": ANY, "epoch": 1}, step=1),
+        call(metrics={"valid_loss_0_epoch": ANY, "valid_loss_1": ANY}, step=1),
     ]
 
     def get_metrics_at_idx(idx):

tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py

@@ -733,16 +733,66 @@ def training_step(self, batch, batch_idx):
     )
     trainer.fit(model)
 
+    # `epoch` is stamped once per step: the epoch-end flush that shares a step with the
+    # last step-level flush no longer re-emits a duplicate `epoch` (see #20902).
     mock_log_metrics.assert_has_calls([
         call(metrics={"foo_step": 0.0, "epoch": 0}, step=0),
         call(metrics={"foo_step": 0.0, "epoch": 0}, step=1),
-        call(metrics={"foo_epoch": 0.0, "epoch": 0}, step=1),
+        call(metrics={"foo_epoch": 0.0}, step=1),
         call(metrics={"foo_step": 0.0, "epoch": 1}, step=2),
         call(metrics={"foo_step": 0.0, "epoch": 1}, step=3),
-        call(metrics={"foo_epoch": 0.0, "epoch": 1}, step=3),
+        call(metrics={"foo_epoch": 0.0}, step=3),
     ])
 
 
+@mock.patch("lightning.pytorch.loggers.TensorBoardLogger.log_metrics")
+def test_log_metrics_no_duplicate_epoch_per_step(mock_log_metrics, tmp_path):
+    """The convenience ``epoch`` metric must be logged at most once per step.
+
+    Regression test for https://github.com/Lightning-AI/pytorch-lightning/issues/20902: the
+    training-epoch-end and validation-end flushes share a step, and stamping ``epoch`` on both
+    produced duplicate ``epoch`` datapoints (rendered as two points per epoch in e.g. MLflow).
+
+    """
+
+    class MyModel(BoringModel):
+        def training_step(self, batch, batch_idx):
+            self.log("train_metric", 1.0, on_step=False, on_epoch=True)
+            return super().training_step(batch, batch_idx)
+
+        def validation_step(self, batch, batch_idx):
+            self.log("val_metric", 2.0, on_step=False, on_epoch=True)
+            return super().validation_step(batch, batch_idx)
+
+    trainer = Trainer(
+        default_root_dir=tmp_path,
+        limit_train_batches=2,
+        limit_val_batches=2,
+        max_epochs=2,
+        log_every_n_steps=1,
+        enable_model_summary=False,
+        enable_checkpointing=False,
+        enable_progress_bar=False,
+        logger=TensorBoardLogger(tmp_path),
+    )
+    trainer.fit(MyModel())
+
+    # count, per step, how many logged batches carried an `epoch` value
+    epoch_counts = collections.Counter()
+    for logged in mock_log_metrics.call_args_list:
+        metrics = logged.kwargs.get("metrics")
+        if metrics is None and logged.args:
+            metrics = logged.args[0]
+        step = logged.kwargs.get("step")
+        if step is None and len(logged.args) > 1:
+            step = logged.args[1]
+        if metrics and "epoch" in metrics:
+            epoch_counts[step] += 1
+
+    assert epoch_counts, "expected `epoch` to be logged at least once"
+    assert all(count == 1 for count in epoch_counts.values()), epoch_counts
+
+
 @mock.patch("lightning.pytorch.loggers.TensorBoardLogger.log_metrics")
 def test_log_on_train_start(mock_log_metrics, tmp_path):
     """Tests that logged metrics on_train_start get reset after the first epoch."""