Lightning-AI / Lightning-AI/pytorch-lightning

metric.compute() hangs when using DDP with multiple GPUs

Open
#20,223 7 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug strategy: ddp ver: 2.4.x
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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.