Lightning-AI / Lightning-AI/pytorch-lightning
self.log: `strategy.reduce` vs `reduce_fx`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
Description & Motivation
self.log used to have separate parameters sync_dist_op and reduce_fx: The reduce_fx callable was used to reduce all of the logged values over an entire epoch, whereas sync_dist_op was used to reduce values across multiple GPUs. Currently it is not possible to set sync_dist_op. Instead, strategy.reduce is used but there is no way to pass the optional reduce_op parameter to this function. I think this means you always get the default, which is 'mean' for DDPStrategy. I think it would be useful for self.log to expose both the between-gpu reduction and the over-epoch reductions.
For example, if one wants to log the maximum value encountered on any GPU over an entire epoch, one would call self.log(..., on_epoch=True, sync_dist=True, sync_dist_op='max', reduce_fx='max'). self.log can then create a partial(strategy.reduce, reduce_op=sync_dist_op) so that the strategy reduce is used, but with the desired reduce_op.
Today, if you do self.log(..., on_epoch=True, sync_dist=True, reduce_fx='max') you will get the maximum over the epoch of the mean value across all GPUs, which is not what is desired, and could be confusing.
Pitch
self.log (and log_dict) should take a sync_dist_op parameter and use it as the reduce_op kwarg for strategy.reduce to improve logging values over multiple GPUs and multiple steps per epoch. The sync_dist_op and reduce_fx parameters should be renamed to make it clear why they are different and what they do.
Alternatives
Use a torchmetrics.Metric instead.
Additional context
No response
cc @borda @carmocca @Blaizzy
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 tracing the self.log and log_dict entry points and the strategy.reduce call, focusing on how reduce_op and reduce_fx currently interact. Compare the current multi-GPU and epoch aggregation behavior with the max-over-GPUs example in the issue. Done means the API clearly supports separate cross-device and over-epoch reductions without the current default-mean ambiguity.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100