From df70affa68d54e58e52533d861d5aac362ecb878 Mon Sep 17 00:00:00 2001 From: Zupeng Wang <71580390+zupengwang@users.noreply.github.com> Date: Sun, 9 Aug 2026 21:28:14 +0800 Subject: [PATCH 1/2] feat: add sync_dist_op to LightningModule logging Signed-off-by: Zupeng Wang <71580390+zupengwang@users.noreply.github.com> --- src/lightning/pytorch/CHANGELOG.md | 2 ++ src/lightning/pytorch/core/module.py | 11 +++++++++++ .../connectors/logger_connector/result.py | 16 ++++++++++++---- .../trainer/logging_/test_train_loop_logging.py | 10 ++++++++++ 4 files changed, 35 insertions(+), 4 deletions(-) diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md index 578c746f1fa5e..6623c1d993091 100644 --- a/src/lightning/pytorch/CHANGELOG.md +++ b/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)) diff --git a/src/lightning/pytorch/core/module.py b/src/lightning/pytorch/core/module.py index 0447c97bf377d..61694b9c069e5 100644 --- a/src/lightning/pytorch/core/module.py +++ b/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,9 @@ 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. Defaults to the operation inferred from + ``reduce_fx`` for backward compatibility. 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 +551,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 +570,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 +598,9 @@ 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. Defaults to the operation inferred from + ``reduce_fx`` for backward compatibility. 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 +637,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 diff --git a/src/lightning/pytorch/trainer/connectors/logger_connector/result.py b/src/lightning/pytorch/trainer/connectors/logger_connector/result.py index 7c9364b5ddfe1..c84d11be4dff7 100644 --- a/src/lightning/pytorch/trainer/connectors/logger_connector/result.py +++ b/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() @@ -371,6 +372,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 @@ -397,7 +399,13 @@ def log( dataloader_idx=self.dataloader_idx, metric_attribute=metric_attribute, ) - 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, + _op=sync_dist_op, + ) # register logged value if it doesn't exist if key not in self: diff --git a/tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py b/tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py index 6916eae68e9c0..ad2f43fa37b9c 100644 --- a/tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py +++ b/tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py @@ -333,6 +333,14 @@ 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", + ) return super().training_step(batch, batch_idx) def validation_step(self, batch, batch_idx): @@ -346,6 +354,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 +391,7 @@ 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["bar"] == fake_result * 3 * devices assert metrics["bar_2"] == fake_result assert metrics["bar_3"] == 2 + int(use_multiple_devices) From b03287f858655842d7621628d8751df48cf021bf Mon Sep 17 00:00:00 2001 From: Zupeng Wang <71580390+zupengwang@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:53:11 +0800 Subject: [PATCH 2/2] Fix explicit rank reductions of logged epoch means --- src/lightning/pytorch/core/module.py | 14 ++++--- .../connectors/logger_connector/result.py | 9 +++- tests/tests_pytorch/core/test_results.py | 42 ++++++++++++++++++- .../logging_/test_train_loop_logging.py | 22 ++++++++++ 4 files changed, 78 insertions(+), 9 deletions(-) diff --git a/src/lightning/pytorch/core/module.py b/src/lightning/pytorch/core/module.py index 61694b9c069e5..3a076d78cd340 100644 --- a/src/lightning/pytorch/core/module.py +++ b/src/lightning/pytorch/core/module.py @@ -425,9 +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. Defaults to the operation inferred from - ``reduce_fx`` for backward compatibility. Has no effect unless ``sync_dist`` is ``True`` and the - strategy is distributed. + 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. @@ -598,9 +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. Defaults to the operation inferred from - ``reduce_fx`` for backward compatibility. Has no effect unless ``sync_dist`` is ``True`` and the - strategy is distributed. + 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. diff --git a/src/lightning/pytorch/trainer/connectors/logger_connector/result.py b/src/lightning/pytorch/trainer/connectors/logger_connector/result.py index c84d11be4dff7..58bf74fe7f2ef 100644 --- a/src/lightning/pytorch/trainer/connectors/logger_connector/result.py +++ b/src/lightning/pytorch/trainer/connectors/logger_connector/result.py @@ -117,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: @@ -148,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 @@ -246,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) @@ -398,13 +403,13 @@ 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, - _op=sync_dist_op, ) # register logged value if it doesn't exist diff --git a/tests/tests_pytorch/core/test_results.py b/tests/tests_pytorch/core/test_results.py index c1d50e8458da6..faa7a01ddd161 100644 --- a/tests/tests_pytorch/core/test_results.py +++ b/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) diff --git a/tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py b/tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py index ad2f43fa37b9c..c9e8a7164196b 100644 --- a/tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py +++ b/tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py @@ -341,6 +341,25 @@ def training_step(self, batch, batch_idx): 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): @@ -392,6 +411,9 @@ def test_logging_sync_dist_true(tmp_path, devices, accelerator): 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)