Lightning-AI / Lightning-AI/pytorch-lightning
sync_dist mean-reduced epoch metrics are inflated when ranks log unequal batch counts (integer batch-size accumulator floors during mean-sync)
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
Epoch-level metrics logged with `self.log(..., on_epoch=True, sync_dist=True)` are **uniformly inflated** whenever ranks accumulate unequal cumulative batch sizes — e.g. an `IterableDataset` whose shards don't divide evenly across ranks, or uneven last batches. Every affected metric is scaled by the same factor `true_denominator / floor(true_denominator / world_size * world_size ... )` — concretely: `(Σbs / W) / floor(Σbs / W)` where `Σbs` is the total logged batch-size mass and `W` the world size.
**Root cause.** `_ResultMetric` accumulates `value` (float) and `cumulated_batch_size` (**int64**, created as `torch.tensor(0)`), and `compute()` mean-syncs *both* through the strategy's `reduce`:
https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/trainer/connectors/logger_connector/result.py#L246-L253
A mean-reduce over an **integer** tensor floors:
- **NCCL**: `_sync_ddp` maps `"mean"` → `ReduceOp.AVG`, and NCCL's AVG on integral dtypes is sum-then-**integer**-division.
- **gloo**: `_sync_ddp` does SUM + `result.copy_(result / world_size)` — the true division result is copied back **into the integer tensor**, truncating:
https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/fabric/utilities/distributed.py#L218-L223
The float numerator keeps its exact fractional value while the integer denominator floors, so the division no longer cancels and the returned mean is inflated. When `Σbs` is divisible by `W` the floor is lossless and everything looks correct — which is why this survives most 2/4/8-rank setups and unit tests, and then silently corrupts larger world sizes.
**Real-world impact (how we found it):** on a 16-rank (4 nodes × 4 GPU) run with a 44-sample validation split logged with `batch_size=1`, every `val` metric was logged ×(44/16)/floor(44/16) = **1.375× too high**, while the identical code on 4 ranks (44/4 = 11, divisible) reproduced single-process values to ≤1e-4. Because *all* mean-reduced metrics scale by the same factor, "higher-is-better" metrics look too good and "lower-is-better" metrics look too bad simultaneously, which is quite misleading during analysis — and any cross-run comparison or val-keyed automation (checkpoint selection across world sizes, early stopping thresholds) silently operates on wrong values.
### Proposed fix
Sync `cumulated_batch_size` in the value's floating dtype so the denominator keeps its fractional part:
```python
cumulated_batch_size = self.meta.sync(self.cumulated_batch_size.to(value.dtype))
```
This keeps the accumulator state itself integer (no checkpoint/state-dict implications) and fixes the reproduction above (returns 7.2). I have a PR ready with this fix plus a 2-rank regression test in `tests_pytorch/core/test_results.py`.
**Related observation** (left out of the PR to keep it scoped): `_sync_ddp` itself floors integer tensors for `"mean"`/`"avg"` on both backends, and `_test_all_reduce` in `tests_fabric/utilities/test_distributed.py` currently *encodes* the floored expectation for integer dtypes (the expected value is cast to the integer dtype before comparison). If maintainers consider integer-mean flooring in `strategy.reduce` itself a bug rather than a contract, I'm happy to follow up separately — it changes the documented in-place semantics (`result is tensor`), so it deserves its own discussion.
### What version are you seeing the problem on?
v2.6, master
### Reproduced in studio
_No response_
### How to reproduce the bug
```python
import torch
from torch.utils.data import DataLoader, IterableDataset
import lightning.pytorch as pl
class RankSizedDataset(IterableDataset):
"""Rank 0 yields 3 samples, rank 1 yields 2 (e.g. sharded data that doesn't divide evenly)."""
def __iter__(self):
rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
values = ([1.0, 2.0, 3.0], [10.0, 20.0])[rank]
return iter([torch.tensor([v]) for v in values])
class Model(pl.LightningModule):
def __init__(self):
super().__init__()
self.layer = torch.nn.Linear(1, 1)
def forward(self, x):
return self.layer(x)
def validation_step(self, batch, batch_idx):
self.log("val_metric", batch.mean(), on_epoch=True, sync_dist=True, batch_size=1)
def configure_optimizers(self):
return torch.optim.SGD(self.parameters(), lr=0.1)
if __name__ == "__main__":
trainer = pl.Trainer(accelerator="cpu", devices=2, strategy="ddp_spawn", logger=False,
enable_progress_bar=False, enable_model_summary=False)
results = trainer.validate(Model(), DataLoader(RankSizedDataset(), batch_size=None))
expected = (1 + 2 + 3 + 10 + 20) / 5 # 7.2, the batch-size-weighted mean over all logged batches
print(f"expected {expected}, got {results[0]['val_metric']}")
```
### Error messages and logs
```
expected 7.2, got 9.0
```
### Environment
Current environment
```
- Lightning: master (35c8970ba) and 2.6.0
- PyTorch: 2.9.1
- Verified on CPU/gloo (script above); NCCL path hit in production on 16 ranks (4 nodes × 4× A10G/L4)
```
### More info
[Fix PR](https://github.com/Lightning-AI/pytorch-lightning/pull/21839)
cc @ethanwharris @lantiga
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 in src/lightning/pytorch/trainer/connectors/logger_connector/result.py at the _ResultMetric.compute() logic, then run the 2-rank regression test in tests_pytorch/core/test_results.py. Done means the uneven-batch reproduction returns the expected 7.2 and the relevant existing tests continue to pass.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, testing
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 25/100