Lightning-AI / Lightning-AI/pytorch-lightning

self.log: `strategy.reduce` vs `reduce_fx`

Open
#17,831 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature help wanted logging pl
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

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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.