Lightning-AI / Lightning-AI/torchmetrics

`MeanAveragePrecision(sync_on_compute=True)` hangs when using `LightningModule` with `ray`'s `TorchTrainer`

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

Nobody has claimed this yet.

bug / fix help wanted waiting on author
Dominant language
Python
Stars
2.5k
Forks
526
Avg merge
6d 11h
Merged PRs (30d)
5

Description

## 🐛 Bug

`torchmetrics.detection.MeanAveragePrecision` (`mAP`) hangs when `sync_on_compute=True` under Distributed Data-Parallel (DDP). By changing this to `sync_on_compute=False`, the job finishes successfully, however, the metric computed on the global rank 0 worker (or any rank worker) is incorrect as no cross-worker state sync occurs before computation.

Calling `metric.sync()` manually also causes the workers to hang; a similar symptom is also reported in #626

### To Reproduce

A full reproducible MRE will take some time to isolate, however, here lie the details of the execution environment.

The job runs distributed supervised 2D object detection training. Each `ray.train.torch.TorchTrainer` worker iterates through a `ray.data.Dataset`, passing data to a `lightning.pytorch.trainer.Trainer` which trains a `lightning.LightningModule` which contains logic to `update()` the `torchmetrics.detection.MeanAveragePrecision` in the `validation_step()`. Finally, in `on_validation_epoch_end()`, we `compute()` the metric, which exhibits the behavior outlined above.

Environment

- TorchMetrics version : `1.6.1`
- torchvision: `0.16.2+cu118`
- OS : `ubuntu`
- torch: `2.1.2+cu118`
- ray: `2.47.1`
- pytorch-lightning: `2.5.0.post0`

### Additional context

- According to this [Stack Overflow comment](https://stackoverflow.com/a/66989632), one of the reasons this can occur is if different workers have a mis-matched number of batches/shards/samples to process - could this be true?
- Switching between the available `coco` backend had no impact
- One way we are considering working around this issue is by setting `sync_on_compute=False`
and using `self.log(sync_dist=True)` in `on_validation_epoch_end()`. This way, the `lightning.LightningModule` handles the aggregation/mean computation on the distributed metric outputs. Mathematically, this is not ideal, but may _temporarily_ unblock work.
- The issue seems to occur when we call `compute --> sync --> gather_all_tensors`

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 by isolating a minimal DDP reproduction around MeanAveragePrecision in validation_step and on_validation_epoch_end, then trace the compute → sync → gather_all_tensors path. Compare sync_on_compute=True with a manual metric.sync() under the reported TorchTrainer and LightningModule setup; done means all workers finish and the metric reflects globally synchronized state.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.