Add sync_dist_op to LightningModule logging

src/lightning/pytorch/CHANGELOG.md

@@ -10,6 +10,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added a `sync_dist_op` argument to `LightningModule.log` and `LightningModule.log_dict` to configure the reduction across devices independently from the epoch reduction ([#17831](https://github.com/Lightning-AI/pytorch-lightning/issues/17831))
+
 - Added `suggest_integrations` flag to `Trainer` to control whether optional integration suggestions (e.g., litmodels, litlogger) are shown in logs ([#21632](https://github.com/Lightning-AI/pytorch-lightning/pull/21632))
 
 - Fixed ``ModelParallelStrategy`` single-file checkpointing when ``torch.compile`` wraps the model so optimizer states no longer raise ``KeyError`` during save ([#21357](https://github.com/Lightning-AI/pytorch-lightning/issues/21357))

src/lightning/pytorch/core/module.py

@@ -37,6 +37,7 @@
 from lightning_utilities.core.apply_func import apply_to_collection
 from lightning_utilities.core.imports import RequirementCache
 from torch import ScriptModule, Tensor
+from torch.distributed import ReduceOp
 from torch.nn import Module
 from torch.optim.optimizer import Optimizer
 from torchmetrics import Metric, MetricCollection
@@ -394,6 +395,7 @@ def log(
         batch_size: Optional[int] = None,
         metric_attribute: Optional[str] = None,
         rank_zero_only: bool = False,
+        sync_dist_op: Optional[Union[ReduceOp, str]] = None,
     ) -> None:
         """Log a key, value pair.
 
@@ -423,6 +425,10 @@ def log(
             sync_dist: if ``True``, reduces the metric across devices. Use with care as this may lead to a significant
                 communication overhead.
             sync_dist_group: the DDP group to sync across.
+            sync_dist_op: the reduction operation for values across devices. For epoch metrics, an explicit operation
+                is applied to the locally reduced values (including local batch-weighted means). If omitted, the
+                operation is inferred from ``reduce_fx`` and epoch means remain globally sample-weighted.
+                Has no effect unless ``sync_dist`` is ``True`` and the strategy is distributed.
             add_dataloader_idx: if ``True``, appends the index of the current dataloader to
                 the name (when using multiple dataloaders). If False, user needs to give unique names for
                 each dataloader to not mix the values.
@@ -546,6 +552,7 @@ def log(
             sync_dist_group=sync_dist_group,
             metric_attribute=metric_attribute,
             rank_zero_only=rank_zero_only,
+            sync_dist_op=sync_dist_op,
         )
 
         trainer._logger_connector._current_fx = self._current_fx_name
@@ -564,6 +571,7 @@ def log_dict(
         add_dataloader_idx: bool = True,
         batch_size: Optional[int] = None,
         rank_zero_only: bool = False,
+        sync_dist_op: Optional[Union[ReduceOp, str]] = None,
     ) -> None:
         """Log a dictionary of values at once.
 
@@ -591,6 +599,10 @@ def log_dict(
             sync_dist: if ``True``, reduces the metric across GPUs/TPUs. Use with care as this may lead to a significant
                 communication overhead.
             sync_dist_group: the ddp group to sync across.
+            sync_dist_op: the reduction operation for values across devices. For epoch metrics, an explicit operation
+                is applied to the locally reduced values (including local batch-weighted means). If omitted, the
+                operation is inferred from ``reduce_fx`` and epoch means remain globally sample-weighted.
+                Has no effect unless ``sync_dist`` is ``True`` and the strategy is distributed.
             add_dataloader_idx: if ``True``, appends the index of the current dataloader to
                 the name (when using multiple). If ``False``, user needs to give unique names for
                 each dataloader to not mix values.
@@ -627,6 +639,7 @@ def log_dict(
                 add_dataloader_idx=add_dataloader_idx,
                 batch_size=batch_size,
                 rank_zero_only=rank_zero_only,
+                sync_dist_op=sync_dist_op,
             )
         return None
 

src/lightning/pytorch/trainer/connectors/logger_connector/result.py

@@ -19,6 +19,7 @@
 import torch
 from lightning_utilities.core.apply_func import apply_to_collection
 from torch import Tensor
+from torch.distributed import ReduceOp
 from torchmetrics import Metric
 from typing_extensions import TypedDict, override
 
@@ -51,7 +52,7 @@ class _Sync:
     fn: Optional[Callable] = None
     _should: bool = False
     rank_zero_only: bool = False
-    _op: Optional[str] = None
+    _op: Optional[Union[ReduceOp, str]] = None
     _group: Optional[Any] = None
 
     def __post_init__(self) -> None:
@@ -68,11 +69,11 @@ def should(self, should: bool) -> None:
         self._generate_sync_fn()
 
     @property
-    def op(self) -> Optional[str]:
+    def op(self) -> Optional[Union[ReduceOp, str]]:
         return self._op
 
     @op.setter
-    def op(self, op: Optional[str]) -> None:
+    def op(self, op: Optional[Union[ReduceOp, str]]) -> None:
         self._op = op
         # `self._fn` needs to be re-generated.
         self._generate_sync_fn()
@@ -116,6 +117,7 @@ class _Metadata:
     add_dataloader_idx: bool = True
     dataloader_idx: Optional[int] = None
     metric_attribute: Optional[str] = None
+    sync_dist_op: Optional[Union[ReduceOp, str]] = None
     _sync: Optional[_Sync] = None
 
     def __post_init__(self) -> None:
@@ -147,7 +149,7 @@ def sync(self) -> _Sync:
     @sync.setter
     def sync(self, sync: _Sync) -> None:
         if sync.op is None:
-            sync.op = self.reduce_fx.__name__
+            sync.op = self.reduce_fx.__name__ if self.sync_dist_op is None else self.sync_dist_op
         self._sync = sync
 
     @property
@@ -245,6 +247,10 @@ def update(self, value: _VALUE, batch_size: int) -> None:
     @override
     def compute(self) -> Tensor:
         if self.is_tensor:
+            if self.meta.is_mean_reduction and self.meta.sync_dist_op is not None:
+                # Apply the explicit rank reduction to local epoch means, not to their
+                # numerators and denominators separately. Keep the default sample-weighted mean below.
+                return self.meta.sync(self.value / self.cumulated_batch_size)
             value = self.meta.sync(self.value.clone())  # `clone` because `sync` is in-place
             if self.meta.is_mean_reduction:
                 cumulated_batch_size = self.meta.sync(self.cumulated_batch_size)
@@ -371,6 +377,7 @@ def log(
         batch_size: Optional[int] = None,
         metric_attribute: Optional[str] = None,
         rank_zero_only: bool = False,
+        sync_dist_op: Optional[Union[ReduceOp, str]] = None,
     ) -> None:
         """See :meth:`~lightning.pytorch.core.LightningModule.log`"""
         # no metrics should be logged with graphs
@@ -396,8 +403,14 @@ def log(
             add_dataloader_idx=add_dataloader_idx,
             dataloader_idx=self.dataloader_idx,
             metric_attribute=metric_attribute,
+            sync_dist_op=sync_dist_op,
+        )
+        meta.sync = _Sync(
+            _should=sync_dist,
+            fn=sync_dist_fn,
+            _group=sync_dist_group,
+            rank_zero_only=rank_zero_only,
         )
-        meta.sync = _Sync(_should=sync_dist, fn=sync_dist_fn, _group=sync_dist_group, rank_zero_only=rank_zero_only)
 
         # register logged value if it doesn't exist
         if key not in self:

tests/tests_pytorch/core/test_results.py

@@ -21,7 +21,7 @@
 from lightning.pytorch.accelerators import CPUAccelerator, CUDAAccelerator, MPSAccelerator
 from lightning.pytorch.strategies import DDPStrategy
 from lightning.pytorch.strategies.launchers import _MultiProcessingLauncher
-from lightning.pytorch.trainer.connectors.logger_connector.result import _Sync
+from lightning.pytorch.trainer.connectors.logger_connector.result import _ResultCollection, _Sync
 from tests_pytorch.helpers.runif import RunIf
 from tests_pytorch.models.test_tpu import wrap_launch_function
 
@@ -54,3 +54,43 @@ def result_reduce_ddp_fn(strategy):
 @RunIf(skip_windows=True)
 def test_result_reduce_ddp():
     spawn_launch(result_reduce_ddp_fn, [torch.device("cpu")] * 2)
+
+
+def result_reduce_epoch_mean_ddp_fn(sync_dist_op, expected, strategy):
+    result = _ResultCollection(training=True)
+    values = (10.0, 14.0) if strategy.global_rank == 0 else (2.0, 4.0)
+    batch_sizes = (1, 1) if strategy.global_rank == 0 else (1, 3)
+    for value, batch_size in zip(values, batch_sizes):
+        result.log(
+            "training_step",
+            "mean",
+            torch.tensor(value),
+            reduce_fx="mean",
+            sync_dist=True,
+            sync_dist_fn=strategy.reduce,
+            sync_dist_op=sync_dist_op,
+            batch_size=batch_size,
+        )
+    # The local means are 12 and 3.5. Without an explicit operation, retain
+    # the global sample-weighted mean (24 + 14) / (2 + 4).
+    torch.testing.assert_close(result["training_step.mean"].compute(), torch.tensor(expected))
+
+
+@pytest.mark.parametrize(
+    ("sync_dist_op", "expected"),
+    [
+        (None, 38 / 6),
+        ("sum", 15.5),
+        ("mean", 7.75),
+        ("avg", 7.75),
+        ("max", 12.0),
+        ("min", 3.5),
+        (dist.ReduceOp.SUM, 15.5),
+        (dist.ReduceOp.MAX, 12.0),
+        (dist.ReduceOp.MIN, 3.5),
+    ],
+)
+@RunIf(skip_windows=True)
+def test_result_reduce_epoch_mean_ddp(sync_dist_op, expected):
+    fn = partial(result_reduce_epoch_mean_ddp_fn, sync_dist_op, expected)
+    spawn_launch(fn, [torch.device("cpu")] * 2)

tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py

@@ -333,6 +333,33 @@ def training_step(self, batch, batch_idx):
         self.log("foo_9", value, on_step=False, on_epoch=True, sync_dist=True, reduce_fx="mean")
         self.log("foo_10", batch_idx, on_step=False, on_epoch=True, sync_dist=True, reduce_fx="max")
         self.log("foo_11", batch_idx + self.rank, on_step=True, on_epoch=True, sync_dist=True, reduce_fx="mean")
+        self.log_dict(
+            {"foo_12": batch_idx + self.rank},
+            on_step=False,
+            on_epoch=True,
+            sync_dist=True,
+            reduce_fx="max",
+            sync_dist_op="sum",
+        )
+        self.log(
+            "foo_13",
+            batch_idx + self.rank,
+            on_step=True,
+            on_epoch=True,
+            sync_dist=True,
+            reduce_fx="mean",
+            sync_dist_op="sum",
+            batch_size=self.rank + 1,
+        )
+        self.log_dict(
+            {"foo_14": batch_idx + self.rank},
+            on_step=False,
+            on_epoch=True,
+            sync_dist=True,
+            reduce_fx="mean",
+            sync_dist_op="mean",
+            batch_size=self.rank + 1,
+        )
         return super().training_step(batch, batch_idx)
 
     def validation_step(self, batch, batch_idx):
@@ -346,6 +373,7 @@ def validation_step(self, batch, batch_idx):
     ("devices", "accelerator"),
     [
         (1, "cpu"),
+        pytest.param(1, "gpu", marks=RunIf(min_cuda_gpus=1)),
         pytest.param(2, "cpu", marks=_xfail_gloo_windows),
         pytest.param(2, "gpu", marks=RunIf(min_cuda_gpus=2)),
     ],
@@ -382,6 +410,10 @@ def test_logging_sync_dist_true(tmp_path, devices, accelerator):
     assert metrics["foo_10"] == 2
     assert metrics["foo_11_step"] == (2 + 3) / 2 if use_multiple_devices else fake_result * 2
     assert metrics["foo_11"] == (0 + 1 + 1 + 2 + 2 + 3) / (devices * 3) if use_multiple_devices else fake_result
+    assert metrics["foo_12"] == (5 if use_multiple_devices else 2)
+    assert metrics["foo_13_step"] == (5 if use_multiple_devices else 2)
+    assert metrics["foo_13_epoch"] == (3 if use_multiple_devices else 1)
+    assert metrics["foo_14"] == (1.5 if use_multiple_devices else 1)
     assert metrics["bar"] == fake_result * 3 * devices
     assert metrics["bar_2"] == fake_result
     assert metrics["bar_3"] == 2 + int(use_multiple_devices)