Skip to content

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

Open
gaurav0107 wants to merge 1 commit into
Lightning-AI:masterfrom
gaurav0107:fix/20902-metrics-get-mapped-twice-to-the-same-epo
Open

gaurav0107 wants to merge 1 commit into
Lightning-AI:masterfrom
gaurav0107:fix/20902-metrics-get-mapped-twice-to-the-same-epo

Conversation

@gaurav0107

Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes #20902

When metrics are logged with on_epoch=True, they show up as two points per epoch in the MLflow UI whenever epoch is used as the x-axis.

Root cause. At an epoch boundary the logger connector flushes several metric groups at the same step — the training-epoch-end flush and the validation-end flush (plus the last train-step flush when on_step metrics are used). Each of these step=None flushes re-stamps the convenience epoch metric via scalar_metrics.setdefault("epoch", ...). Loggers that keep per-call history (MLflow logs one Metric entry per log_metrics call) therefore record two epoch datapoints for every real epoch, so the epoch series has twice as many points as the other metric series and the x-axis mapping is off.

Fix. 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. All real metrics and all step values are unchanged; step-indexed loggers (TensorBoard, CSV) are unaffected because they already overwrite / merge by step.

This is a minimal, backward-compatible change contained to _LoggerConnector.log_metrics. Two existing tests encoded the old (duplicated) behavior and are updated to assert the corrected single-epoch-per-step behavior, and a focused regression test is added.

Before submitting
  • Was this discussed/agreed via a GitHub issue? Yes — Metrics get mapped twice to the same epoch in MLflow logger #20902 (labeled bug).
  • Did you read the contributor guideline, Pull Request section?
  • 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? Not needed (behavior fix).
  • Did you write any new necessary tests? Added test_log_metrics_no_duplicate_epoch_per_step and updated existing epoch/step assertions.
  • Did you verify new and existing tests pass locally with your changes? (see note below)
  • Did you list all the breaking changes introduced by this pull request? None — epoch is still logged once per step.
  • Did you update the CHANGELOG?

Note: the sandbox used to prepare this change does not have a full PyTorch/Lightning runtime, so the pytest suite is validated by CI. Locally verified: ruff check, ruff format --check, and py_compile on all touched files pass. The updated assertions follow directly from the connector's flush order (the first flush at each step keeps epoch).

PR review

Anyone in the community is welcome to review the PR.

@gaurav0107
gaurav0107 marked this pull request as ready for review July 1, 2026 21:49
@gaurav0107
gaurav0107 force-pushed the fix/20902-metrics-get-mapped-twice-to-the-same-epo branch from 1440fdc to 2a5ff23 Compare July 2, 2026 13:59
@codecov-commenter

codecov-commenter commented Jul 2, 2026 •

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79%. Comparing base (fbdf042) to head (d47e78b).
⚠️ Report is 15 commits behind head on master.
✅ All tests successful. No failed tests found.
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

❗ There is a different number of reports uploaded between BASE (fbdf042) and HEAD (d47e78b). Click for more details.

HEAD has 2323 uploads less than BASE
Flag BASE (fbdf042) HEAD (d47e78b)
cpu 583 59
python 29 3
lightning_fabric 222 0
pytest 297 0
python3.12 205 20
python3.10 30 3
lightning 149 15
python3.13 170 18
python3.12.7 89 9
python3.11 60 6
pytorch2.2.2 15 3
pytest-full 286 59
pytorch2.3 15 3
pytorch2.1 29 6
pytorch_lightning 212 44
pytorch2.9 29 6
pytorch2.13 29 6
pytorch2.4.1 14 3
pytorch2.5.1 15 3
pytorch2.10 27 5
pytorch2.12 27 6
pytorch2.11 28 6
pytorch2.8 29 6
pytorch2.7 14 3
pytorch2.6 15 3
Additional details and impacted files
@@            Coverage Diff            @@
##           master   #21796     +/-   ##
=========================================
- Coverage      87%      79%     -8%     
=========================================
  Files         270      267      -3     
  Lines       24074    24021     -53     
=========================================
- Hits        20906    18873   -2033     
- Misses       3168     5148   +1980     
…0902)

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.
@gaurav0107
gaurav0107 force-pushed the fix/20902-metrics-get-mapped-twice-to-the-same-epo branch from 2a5ff23 to d47e78b Compare July 16, 2026 18:58

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

2 participants