WandbLogger: add more flexibility with logging step

[borisdayma]

🚀 Feature

Let WandbLogger log using:

• the trainer step by default (current behavior)
• an auto-incremented step, independent from the trainer

Motivation

The current default of using trainer step is a good default and seems to work fine for most people, automatically associated correctly training and validation metrics at the right step.

However, some users have a few issues with custom logging:

• wandb logger problem with on_step log on validation #4980 wants to see a chart of validation metrics over a single batch
• Logging validation metrics with val_check_interval #5070 tries to log intermediate validation metrics, which uses a different step
• "Step must only increase in log calls" when adding W&B logger after some training wandb/wandb#1626
• WandB-Logger drops all the logged values in training step for PyTorch Lightning wandb/wandb#1507

#5050 solves some of these issues by suggesting using wandb.log(my_dict, commit=False) when logging outside of the "regular" workflow (which will automatically associate metrics to last used step) but not all of them.
In some cases it may be hard to ensure the associated step is always increasing.

Pitch

We will be able to have a new parameter in WandbLogger (TBD) such as sync_step=True:

• by default True, we will have the current behavior which syncs wandb.step with the trainer step.
• when set to False, the step will auto-increment at each logging (by relying on wandb.log() default behaviour) and the trainer step will just be added to logged metrics when available

In those cases (set to False), we will still be able to manually select the trainer step as x-axis in the W&B dashboard (global graphs or any specific graph) so this feature will add more flexibility in logging.

It will not be the default behavior mainly for people not familiar with W&B and that expect to see validation and training metrics aligned by default (see #4113).

[awaelchli]

I like it. I believe it currently is also possible to self.log("step", ...) and the logger will use this to log the metrics.

[Shreeyak]

@awaelchli Could you please elaborate on what you meant by:

is also possible to self.log("step", ...) and the logger will use this to log the metrics.

I didn't quite understand how to use your suggestion to log values per validation step (if that is what you meant).

[Shreeyak]

when set to False, the step will auto-increment at each logging

@borisdayma What about the cases when multiple log statements are present and they should be grouped into the same step? Would it make sense to add the commit argument from wandb api as well?

[awaelchli]

@Shreeyak sorry, it used to work that way, or at least for TensorBoard. We could think of adding it and make it consistent across all loggers, but it should be used with care because most loggers I know don't support logging to the "past", for good reason :)

What about the cases when multiple log statements are present and they should be grouped into the same step? Would it make sense to add the commit argument from wandb api as well?

not necessary, internally we call wandb.log with all logs combined. commit=False is only necessary when user invokes the wandb api directly.

[borisdayma]

What about the cases when multiple log statements are present and they should be grouped into the same step? Would it make sense to add the commit argument from wandb api as well?

So by default I would always keep the current behavior.
What will happen with the new (optional) version (sync_step=False) will be that each log call will be associated to a different wandb.step (automatically incrementing).
As mentioned by @awaelchli, those metrics may be grouped into a dict in the same log call but the wandb.step will be independent from the trainer's step.

However, even when the metrics are not logged in the same call, or if you do other calls to wandb.log (for example to log media, images, audio…), you will be able to adjust in your interface your graph and change the x-axis because the trainer step will be logged a a metric.

You can change the x-axis for all graphs or on a per-graph basis, as long as the step is always increasing.

The main advantage of this method is actually for people who do really custom logging where the current behavior does not fit their use:

• logging per-batch metrics during validation
• bugs where step increases during their custom calls, causing some issues for logging back in past -> this is often due to calls to wandb.log without commit=False in their validation loop but there may be other situations which I've not been able to observe or understand myself

Feel free to bring up any concerns or issues.

[Shreeyak]

Thanks! So if using sync_step=False, users would have to group the metrics they want to log in a single step? Example: if my validation loop has metric_1 and metric_2, I'd log them together similar to this: wandb.log({'metric_1': 0.9, 'metric_2': 0.1})?
And otherwise we'd increment wandb step for each call?

I'm thinking of use-cases where the validation loop takes hours and it'd be useful to see some metrics along the way.

You bring up a good point of being able to adjust the x-axis in wandb though.

[borisdayma]

@Shreeyak You can test PR #5351 which should just be missing the tests (I'm experimenting with it manually).

My ideal goal would be that everything works perfectly for everyone with sync_step=True. I noticed that when it does not it is usually due to users doing manual calls to wandb.log (often for valid reasons), which gets solved by adding wandb.log(commit=False). However there has been some instances where users had still this issue and I've not been able to find out why due to lack of reproducible examples.

If you use sync_step=False, you can still log with regular PL behavior self.log('my_metric', my_value) and PL will log everything at a time (for one batch) or you can do manual calls to wandb.log. In that case, it does not matter if you group or not your metrics, just know that wandb will increment its internal step at each logging.
The advantage of using self.log will be that PL will also add a metric trainer_step which you can later use as your x-axis and which will give you the same graphs that if you used sync_step=True.

If you have a chance, I'd love to see you test current PR and see if that works for you.