Lightning-AI / Lightning-AI/pytorch-lightning

Allow to specify the value of the `trainer/global_step` metric on W&B during validation

Open
#14,892 6 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

logger: wandb question won't fix
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

## 🚀 Feature

Allow to specify the value of the `trainer/global_step` metric on W&B even during validation, as it is possible during training.

### Motivation

The default value for the `trainer/global_step` metric during validation is defined as the "total validation step":

https://github.com/Lightning-AI/lightning/blob/e7ad151c311890c931aa5ff3df3672fd21fd351c/src/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py#L152

whereas during training the default value is `trainer.fit_loop.epoch_loop._batches_that_stepped`, even if there are some inconsistencies (see #13752).

The fact that the behaviour of `trainer/global_step` is different when switching from training to validation, generates graphs like this on W&B:

image

as well as the fact that W&B complains when trying to use `trainer/global_step` as `x` axis because it's not monotonically increasing:

image

For training it's possible to set a custom value for `training/global_step` by using `self.log("step", ...)` during the training step. For validation, there's no such possibility.

I understand that if one logs using `on_step=True` during validation, it's important to keep track of the different validation steps. But if one logs using `on_step=False` and `on_epoch=True`, it would be useful to link the val metric to the current training step instead of the total validation step.

### Pitch

One can set a custom value for `training/global_step` even during validation.

### Alternatives

Another possibility could be to log a different metric for validation, e.g. `trainer/validation_global_step` and to use that as default x axis as it is done with `trainer/global_step` here:

https://github.com/Lightning-AI/lightning/blob/6df6deac8fdf0c44b509d531441b4f7e446de6b4/src/pytorch_lightning/loggers/wandb.py#L396-L397

cc @awaelchli @morganmcg1 @borisdayma @scottire @parambharat @manangoel99

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by comparing validation step handling in src/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py with training's fit_loop.epoch_loop._batches_that_stepped, then inspect the W&B logger entry point in src/pytorch_lightning/loggers/wandb.py. Done means validation logging can use a specified global-step value without breaking per-step validation tracking or W&B's monotonic-axis expectations.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, observability
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.