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

fix: log the convenience epoch metric once per step (#20902)#21796
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
…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 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 (4819088) to head (2a5ff23).
✅ 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 (4819088) and HEAD (2a5ff23). Click for more details.

HEAD has 168 uploads less than BASE
Flag BASE (4819088) HEAD (2a5ff23)
python3.10 6 3
cpu 84 42
lightning 30 15
pytest 42 0
python 6 3
lightning_fabric 27 0
python3.12 24 12
python3.12.7 18 9
python3.11 12 6
python3.13 18 9
Additional details and impacted files
@@            Coverage Diff            @@
##           master   #21796     +/-   ##
=========================================
- Coverage      87%      79%     -8%     
=========================================
  Files         270      267      -3     
  Lines       24010    23957     -53     
=========================================
- Hits        20787    18839   -1948     
- Misses       3223     5118   +1895     

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Metrics get mapped twice to the same epoch in MLflow logger

2 participants