Lightning-AI / Lightning-AI/torchmetrics
`MeanAveragePrecision(sync_on_compute=True)` hangs when using `LightningModule` with `ray`'s `TorchTrainer`
Nobody has claimed this yet.
- 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
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 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