Lightning-AI / Lightning-AI/torchmetrics
`_sync_dist` hangs for `dist_reduce_fx=None` list states when ranks have different list lengths
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 2.5k
- Forks
- 526
- Avg merge
- 6d 11h
- Merged PRs (30d)
- 5
Description
## Description
`_sync_dist` deadlocks when a metric uses `dist_reduce_fx=None` list states and different ranks have different numbers of entries in those lists. This affects `MeanAveragePrecision`, retrieval metrics, and any custom metric that declares list states with `dist_reduce_fx=None`.
The same class of bug was fixed for `dist_reduce_fx="cat"` in #2468 (v1.3.x), but `None`-reduction states were not covered and require a fundamentally different fix.
### To Reproduce
Provided are two code samples to reproduce, one that is hypothetical where I construct a "None reduction list-based metric", and another that is more realistic using the `MeanAveragePrecision` metric.
Realistic example
```python
"""
Minimal reproduction: MeanAveragePrecision hangs on compute() with
sync_on_compute=True when one rank has empty states (no update() called).
MeanAveragePrecision uses dist_reduce_fx=None for all its list states.
The empty-list corner case was fixed for dist_reduce_fx="cat" in
https://github.com/Lightning-AI/torchmetrics/pull/2468 (v1.3.x),
but dist_reduce_fx=None was not covered.
Run: torchrun --nproc_per_node=2 repro_torchmetrics_hang.py
Expected: hangs at metric.compute() on rank 1 (waiting for all_gather)
"""
import signal
import torch
import torch.distributed as dist
from torchmetrics.detection import MeanAveragePrecision
# Auto-kill after 30s so CI doesn't wait forever
signal.alarm(30)
def make_sample_prediction(device):
"""Create a minimal valid prediction + target for MeanAveragePrecision."""
preds = [
{
"boxes": torch.tensor([[10.0, 10.0, 50.0, 50.0]], device=device),
"scores": torch.tensor([0.9], device=device),
"labels": torch.tensor([1], device=device),
}
]
targets = [
{
"boxes": torch.tensor([[10.0, 10.0, 50.0, 50.0]], device=device),
"labels": torch.tensor([1], device=device),
}
]
return preds, targets
def main():
dist.init_process_group("nccl")
rank = dist.get_rank()
device = torch.device(f"cuda:{rank}")
torch.cuda.set_device(device)
metric = MeanAveragePrecision(sync_on_compute=True).to(device)
if rank == 1:
preds, targets = make_sample_prediction(device)
metric.update(preds, targets)
print(f"[Rank {rank}] Called update() with 1 sample")
else:
print(f"[Rank {rank}] Skipping update() -- states remain as empty lists")
print(f"[Rank {rank}] Calling compute()...")
# Rank 0: all 9 list states are [] -> apply_to_collection finds 0 tensors -> 0 all_gather calls
# Rank 1: list states are populated -> apply_to_collection finds N tensors -> N all_gather calls
# Rank 1 blocks at the first all_gather waiting for rank 0 -> HANG
result = metric.compute()
print(f"[Rank {rank}] compute() returned: {result}") # rank 1 never reaches this
dist.destroy_process_group()
if __name__ == "__main__":
main()
```
Run the example with `torchrun --nproc_per_node=2 repro.py`. The `MeanAveragePrecision.compute()` deadlocks.
Theoretical example
```python
"""
Minimal reproduction: dist_reduce_fx=None list states hang on compute()
when one rank has empty states.
The empty-list corner case was fixed for dist_reduce_fx="cat" in
https://github.com/Lightning-AI/torchmetrics/pull/2468 (v1.3.x),
but dist_reduce_fx=None was not covered. This script demonstrates
the hang using a trivial custom metric.
Run: torchrun --nproc_per_node=2 repro_torchmetrics_hang.py
Expected: hangs at metric.compute() on rank 1 (waiting for all_gather)
"""
import os
import signal
import torch
import torch.distributed as dist
import torchmetrics
# Auto-kill after 30s so CI doesn't wait forever
signal.alarm(30)
class NoneReductionListMetric(torchmetrics.Metric):
"""Metric with a list state using dist_reduce_fx=None (same as MeanAveragePrecision)."""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.add_state("items", default=[], dist_reduce_fx=None)
def update(self, x: torch.Tensor):
self.items.append(x)
def compute(self):
if len(self.items) == 0:
return torch.tensor(0.0, device=self.device)
return torch.stack(self.items).sum()
class CatReductionListMetric(torchmetrics.Metric):
"""Metric with a list state using dist_reduce_fx='cat' (protected by the existing fix)."""
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.add_state("items", default=[], dist_reduce_fx="cat")
def update(self, x: torch.Tensor):
self.items.append(x)
def compute(self):
if len(self.items) == 0:
return torch.tensor(0.0, device=self.device)
print(f"{self.items=}")
return self.items.sum()
def main():
dist.init_process_group("nccl")
rank = dist.get_rank()
device = torch.device(f"cuda:{rank}")
torch.cuda.set_device(device)
print(f"[Rank {rank}] Testing dist_reduce_fx='cat' (should NOT hang)...")
metric_cat = CatReductionListMetric(sync_on_compute=True).to(device)
if rank == 1:
metric_cat.update(torch.tensor([1.0, 2.0], device=device))
metric_cat.update(torch.tensor([1.0, 2.0], device=device))
# Rank 0: no update -> empty list state
result = metric_cat.compute() # Should succeed (protected by existing fix)
print(f"[Rank {rank}] dist_reduce_fx='cat' OK: {result}")
dist.barrier()
print(f"[Rank {rank}] Testing dist_reduce_fx=None (will HANG)...")
metric_none = torchmetrics.MetricCollection({"m": NoneReductionListMetric(sync_on_compute=True).to(device)})
if rank == 1:
metric_none.update(torch.tensor([1.0, 2.0], device=device))
# Rank 0: no update -> empty list state
# This will hang: rank 1 calls all_gather inside _sync_dist,
# rank 0 finds no tensors in its empty list and skips all_gather.
result = metric_none.compute() # <-- HANGS HERE
print(f"[Rank {rank}] dist_reduce_fx=None OK: {result}") # never reached
dist.destroy_process_group()
if __name__ == "__main__":
main()
```
Run the example with `torchrun --nproc_per_node=2 repro.py`. The `cat`-reduction metric completes; the `None`-reduction metric deadlocks.
Environment
- TorchMetrics version (if build from source, add commit SHA): 1.9.0
- Python & PyTorch Version (e.g., 1.0): Python: 3.10.19 | packaged by conda-forge | (main, Oct 22 2025, 22:29:10) [GCC 14.3.0]
- PyTorch: 2.10.0+cu128
- CUDA available: True
- CUDA version: 12.8
- GPU: NVIDIA A10G
- OS: Linux-6.5.0-1024-aws-x86_64-with-glibc2.35
## Root cause
In `_sync_dist`, `apply_to_collection` maps `gather_all_tensors` (which calls `torch.distributed.all_gather`) over every tensor found in each state list. The number of `all_gather` calls equals the number of tensors in the list. When two ranks have different list lengths, they issue different numbers of collectives and deadlock.
```python
# metric.py _sync_dist (simplified)
output_dict = apply_to_collection(input_dict, Tensor, gather_all_tensors, ...)
# For a list [t1, t2, ...], this calls gather_all_tensors once per element.
# Rank with 0 elements: 0 calls. Rank with N elements: N calls. → deadlock.
```
## Why the existing `dim_zero_cat` fix does not generalize to `None`
The fix in #2468 works for `dim_zero_cat` because of two properties that `None` does not share:
### 1. Pre-concatenation normalizes list length to 1
Lines 506-507 of `metric.py`:
```python
if reduction_fn == dim_zero_cat and isinstance(input_dict[attr], list) and len(input_dict[attr]) > 1:
input_dict[attr] = [dim_zero_cat(input_dict[attr])]
```
This collapses any length-N list into a single-element list before gathering. Combined with the empty-list placeholder (lines 510-516), every rank always has exactly 1 list element → 1 `all_gather` call → symmetric.
For `None`, there is no pre-concatenation. Lists can have any length, and concatenating them would destroy the per-entry structure that `compute()` depends on (e.g., `MeanAveragePrecision` needs each list entry to represent one image's detections).
### 2. Reduction eliminates the placeholder
For `dim_zero_cat`, the placeholder `torch.tensor([])` is eliminated by the reduction step:
```python
torch.cat([torch.tensor([]), real_data]) → real_data # empty contributes nothing
```
For `None`, the reduction is identity — the placeholder would survive into `compute()` as a real list entry with potentially wrong tensor shape.
### 3. ndim mismatch
The placeholder is a 1-D tensor `torch.tensor([])`, but `None`-reduction states often contain multi-dimensional tensors (e.g., `detection_box` entries are shape `(num_boxes, 4)`, ndim=2). `gather_all_tensors` exchanges shapes via `all_gather` as its first step, which requires all ranks to send tensors with the same `ndim`. A 1-D placeholder paired with a 2-D real tensor causes a shape mismatch in this size-exchange step.
## Impact
This affects any metric with `dist_reduce_fx=None` list states when used with `sync_on_compute=True` and uneven data across ranks. Known affected metrics:
- **`MeanAveragePrecision`** — all 9 states use `dist_reduce_fx=None`
- **Retrieval metrics** — use `dist_reduce_fx=None` list states
- **Custom metrics** following the pattern `self.add_state("foo", default=[], dist_reduce_fx=None)`
Common scenarios that trigger uneven list lengths:
- `drop_last=False` in the dataloader (last batch goes to some ranks but not others)
- Filtered/conditional evaluation where some batches have no valid targets on some ranks
- Intentionally skipping `update()` on certain ranks
## Possible fix directions
A proper fix would need to equalize list lengths across ranks before `apply_to_collection`. Unlike `dim_zero_cat`, the entries cannot be concatenated (per-entry identity matters for `compute()`), so the fix would need to:
1. Exchange list lengths across ranks (e.g., `all_reduce` with `MAX`)
2. Pad shorter lists with ndim-matched empty tensors (shape `(0, ...)` matching the real tensors' ndim)
3. Proceed with `apply_to_collection` as normal — all ranks now have the same number of elements
This requires knowing the expected `ndim` per state. One approach: if any rank has at least one real tensor, broadcast its `ndim` during the length exchange. If all ranks are empty, no gathering is needed.
Alternatively, a simpler mitigation would be to detect the mismatch and raise an informative error instead of silently deadlocking:
```python
if reduction_fn is None and isinstance(input_dict[attr], list):
local_len = torch.tensor(len(input_dict[attr]), device=self.device)
all_lens = [torch.zeros_like(local_len) for _ in range(world_size)]
dist.all_gather(all_lens, local_len)
if len(set(l.item() for l in all_lens)) > 1:
raise RuntimeError(
f"Metric state '{attr}' has dist_reduce_fx=None but different list lengths "
f"across ranks: {[l.item() for l in all_lens]}. This will deadlock in "
f"apply_to_collection. Ensure all ranks call update() the same number of "
f"times, or use sync_on_compute=False."
)
```
## Application-level workaround
Until this is fixed upstream, the workaround is to **always call `update()` with zero-length tensors** when a rank has no real data, rather than skipping `update()`:
```python
# Instead of skipping update():
if no_valid_data:
empty_preds = [{"boxes": torch.zeros((0, 4), device=device),
"scores": torch.zeros(0, device=device),
"labels": torch.zeros(0, dtype=torch.long, device=device)}]
empty_targets = [{"boxes": torch.zeros((0, 4), device=device),
"labels": torch.zeros(0, dtype=torch.long, device=device)}]
metric.update(empty_preds, empty_targets)
```
Zero-length tensors (shape `(0, 4)` not `(1, 4)`) pass through `_input_validator` and `_get_safe_item_values` cleanly, keep list lengths synchronized across ranks, have matching `ndim` for `gather_all_tensors`, and contribute nothing to the metric computation (0 boxes = 0 TP/FP/FN).
## Related issues
- #2463 — original report of the empty-list hang for `dim_zero_cat` (fixed in #2468)
- #2170 — imbalanced metric states with `drop_last=False` (same underlying issue, different trigger)
- #2481 — retrieval metrics GPU memory leak with `dist_reduce_fx=None` list states
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 metric.py at _sync_dist and trace how apply_to_collection invokes gather_all_tensors for list states. Reproduce the uneven-list case with the provided torchrun examples, including MeanAveragePrecision and a custom dist_reduce_fx=None metric. Done means distributed compute no longer deadlocks when ranks have different list lengths, while preserving per-entry state behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100