Lightning-AI / Lightning-AI/torchmetrics

`_sync_dist` hangs for `dist_reduce_fx=None` list states when ranks have different list lengths

Open
#3,336 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug / fix help wanted
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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.