Lightning-AI / Lightning-AI/pytorch-lightning
metric.compute() hangs when using DDP with multiple GPUs
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
I'm using the default Accuracy metric (though it appears to be true for any metric), and calling metric.compute() hangs after the first epoch and never resolves (ran it overnight, never progressed). It seems, as per some print() statements, that the issue is only with metric computation after training epoch ends, not after validation epoch ends. Issue does not happen for when using only 1 gpu or cpu. It is also agnostic of how large the dataset is, I tried with a dataset having only the first 2 batches and got the same result. I see there's another relevant issue (#5930) from 3 years ago, but has no solution (just says to update version and make a new issue).
### What version are you seeing the problem on?
v2.4
### How to reproduce the bug
In Model(pl.LightningModule).__init__(self, splits, #more args):
```python
# ...
# without "_metrics" suffix, keys conflicts with nn.ModuleDict, Adding split prefix as well, but not sure if that is needed
self.split_metrics = nn.ModuleDict({f'{split}_metrics': nn.ModuleDict({f'{split}_{name.replace(".","")}': metric for name, metric in {
'acc': torchmetrics.Accuracy(task='binary'),
# ... more metrics
}.items()}) for split in splits})
# ...
```
(Note: I considered using a metric collection, but some of my metrics need different inputs and I couldn't figure out how to account for that)
In Model.(pl.LightningModule)._step(self, batch, batch_idx, *, split, **kwargs):
```python
# ...
for name, metric in self.split_metrics[f'{split}_metrics'].items():
metric.update(y_pred_prob if needs_probability(name) else y_pred_label, y)
# ...
```
Relevant overloads in Model(pl.LightningModule):
```python
def training_step(self, *args, **kwargs):
return self._step(*args, **kwargs, split='train')
def validation_step(self, *args, **kwargs):
return self._step(*args, **kwargs, split='val')
def on_train_epoch_end(self, *args, **kwargs):
self._on_epoch_end(*args, **kwargs, split='train')
def on_validation_epoch_end(self, *args, **kwargs):
self._on_epoch_end(*args, **kwargs, split='val')
```
### Error messages and logs
I put some print statements that show it hangs for train_acc. This is with 2 GPUs.
```
Sanity Checking DataLoader 0: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 2.80it/s][0] VALIDATION EPOCH END
[0] ENTERING _on_epoch_end :: val
[0] log dict initialized :: val
[0] now in loop for val_acc :: val
[1] VALIDATION EPOCH END
[1] ENTERING _on_epoch_end :: val
[1] log dict initialized :: val
[1] now in loop for val_acc :: val
[1] computed val_acc :: val
[1] logged val_acc :: val
[1] reset val_acc :: val
[0] computed val_acc :: val
[0] logged val_acc :: val
[0] reset val_acc :: val
Epoch 0: 0%| | 0/2 [00:00
Current environment
```
#- PyTorch Lightning Version (e.g., 2.4.0): 2.4.0
#- PyTorch Version (e.g., 2.4): 2.4.0
#- TorchMetrics Version: 1.4.1
#- Python version (e.g., 3.12): 3.11.9
#- OS (e.g., Linux): Linux
#- CUDA/cuDNN version: 12.1
#- GPU models and configuration: 8 x Tesla V100-SXM2-16GB
#- How you installed Lightning(`conda`, `pip`, source): pip
```
### More info
_No response_
cc @justusschock @lantiga
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 with the metric.compute() call in _on_epoch_end, comparing the train path from on_train_epoch_end with the working validation path from on_validation_epoch_end. Reproduce the hang using the supplied 2-GPU setup and inspect the train_acc update, compute, log, and reset sequence. Done means metric computation completes after a training epoch under DDP without hanging.
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
- Mostly clear
- Newbie friendliness
- 35/100