Lightning-AI / Lightning-AI/pytorch-lightning

torch.randn() + DDP + GANs are easy to get wrong with lightning

Open
#16,166 3 comments 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

discussion feature help wanted reproducibility
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

If you run the GAN example model https://github.com/Lightning-AI/lightning/blob/master/examples/pl_domain_templates/generative_adversarial_net.py with DDP you are effectively training with only a single GPU because they are all sampling the same latent vector. The offending code is https://github.com/Lightning-AI/lightning/blob/3ff3ec3fdef92fa2f187f06eca41bf08dcc4eb19/examples/pl_domain_templates/generative_adversarial_net.py#L155 to fix this I do something like this in every GAN model I create.
```python
def _augmentation_seed(self):
if not hasattr(self, 'seeded') or not self.seeded:
seeds = torch.randint(0, 2**32 - 1, (self.trainer.world_size, ))
pl.seed_everything(seeds[self.trainer.global_rank], True)
self.seeded = True

def _sample_latent(self, imgs):
self._augmentation_seed()
return torch.randn(imgs.shape[0], self.hparams.latent_dim)

def training_step(self, batch, batch_idx, optimizer_idx):
imgs, _ = batch

# sample noise
z = self._sample_latent(imgs)
z = z.type_as(imgs)
```
It would be nice if lightning could either 1 warn the user about this or 2 (I think this would be better) after all models have been initialized, something like _augmentation_seed is called internally to the Lightning module. I call _augmentation_seed in training_step because I am not 100% sure when all models have been initialized. I also want to point out that this is not only restricted to GANs; it would also affect anyone doing augmentation in the training loop.

### How to reproduce the bug

_No response_

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current environment

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 1.10):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):
```

### More info

_No response_

cc @borda @awaelchli

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 examples/pl_domain_templates/generative_adversarial_net.py at the referenced latent-sampling code, then inspect how DDP initializes model replicas and random state. Reproduce the GAN example under DDP if possible and verify that each rank samples independently, or that the framework clearly warns about shared randomness.

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.