Lightning-AI / Lightning-AI/pytorch-lightning
Allow to specify the value of the `trainer/global_step` metric on W&B during validation
Nobody has claimed this yet.
- 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:

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:

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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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