Lightning-AI / Lightning-AI/pytorch-lightning

DDP + static graph can result in garbage data returned by `all_gather`

Open
#18,872 2 comments 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

3rd party bug repro needed ver: 2.0.x
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

When I use self.all_gather in LightningModule with strategies.DDPStrategy(static_graph=True) for multi-node inference,
the returned values are partially corrupted.

### What version are you seeing the problem on?

v2.0

### How to reproduce the bug

```python
class ModelWrapper(pl.LightningModule):
...
def on_validation_epoch_end(self) -> None:
super().on_validation_epoch_end()

test_tensor = torch.tensor([self.global_rank]).to(current_device)
self.trainer.strategy.barrier()
test_all = self.all_gather(test_tensor)
self.trainer.strategy.barrier()
print(test_all)

```

And it is called by
```python
# set up a datamodule
...

trainer = pl.Trainer(
accelerator="gpu",
devices=args.devices,
num_nodes=args.num_nodes,
strategy=strategies.DDPStrategy(static_graph=True),
)
trainer.validate(model, datamodule=datamodule)
```

### Error messages and logs

It should return results as:
torch.tensor([0, 1, 2, ..., world_size-1])

But while most of values are right, a few comes back with corrupted data with very large numbers, like
torch.tensor([0, 1, 2, 3, 913478191043, 5, ..., world_size -1])

### Environment

Current environment

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow): LightningModule
#- PyTorch Lightning Version (e.g., 1.5.0): 2.0.3
#- PyTorch Version (e.g., 2.0): 1.12.1
#- Python version (e.g., 3.9): 3.9
#- OS (e.g., Linux): Linux
#- CUDA/cuDNN version: 11.6
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source): conda
```

### Addtional Information
I do notice the warning that trainer.validate should not be called with DDPStrategy, which makes LightningModule copies some datapoints for the last round of validation. Actually, this is exactly why I use all_gather during validation - to implement a drop-last validation.
It looks like it's caused by failure to block all processes during all_gather. I've tried to investigate why this happens myself, but I can't find any clues.

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

No source files or tests are named. Start by running the provided multi-node validation reproduction with DDPStrategy(static_graph=True), then trace LightningModule.all_gather and the surrounding barrier behavior. Done means all_gather returns the expected rank values without corrupted entries.

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.