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)

Open
#21,859 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug logger ver: 2.6.x ver: 2.7.x
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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.