Lightning-AI / Lightning-AI/pytorch-lightning

Fabric: Incorrect `num_replicas` (ddp/fsdp) when number of GPUs on each node is different

Open
#19,898 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug fabric strategy: ddp
Dominant language
Python
Stars
31.4k
Forks
3.8k
Avg merge
6d 7h
Merged PRs (30d)
6

Description

### Bug description

When running multi-node/multi-GPU training with different number of GPUs on each node, `Fabric` `ddp` and `fsdp` will have an incorrect `num_replicas` in `distributed_sampler_kwargs`: Currently `num_replicas` is set to be `num_gpus * num_nodes`, instead of simply `world_size`.

To reproduce the bug: run fabric on two nodes, one with 2 GPUs and another with only one (the `global_rank` on the second node should be `2`).
In that case, the `num_replicas` on will be different on the two nodes: while on the node with two GPUs it will be 4, on the node with one GPU it will be 2.

Why not setting `num_replicas` to simply be `world_size`?

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

v2.2

### How to reproduce the bug

```python
Run lightning Fabric on three GPUs on two different nodes

from lightning.fabric import Fabric
from torchvision.datasets import MNIST
from torchvision import transforms as tvt
from torch.utils.data import DataLoader

# num_devices is either 2 or 1 depending on the node
num_devices =
fabric = Fabric(accelerator='cuda', strategy='ddp',
devices=num_devices, num_nodes=2)
fabric.launch()

# pick the simplest Dataset you want
train_data = DataLoader(MNIST(root='.', train=True, transform=tvt.ToTensor()), batch_size=3, num_workers=2)

# this will fail on the node with single GPU
train_loader = fabric.setup_dataloaders(train_data)
```

### Error messages and logs

```
Traceback (most recent call last):
...
File ".../main_linprobe.py", line 225, in main
data_loader_train = fabric.setup_dataloaders(data_loader_train)
File ".../site-packages/lightning/fabric/fabric.py", line 376, in setup_dataloaders
dataloaders = [
File ".../site-packages/lightning/fabric/fabric.py", line 377, in
self._setup_dataloader(
File ".../site-packages/lightning/fabric/fabric.py", line 404, in _setup_dataloader
sampler = self._get_distributed_sampler(dataloader, **self._strategy.distributed_sampler_kwargs)
File ".../site-packages/lightning/fabric/fabric.py", line 1005, in _get_distributed_sampler
return DistributedSampler(dataloader.dataset, **kwargs)
File ".../site-packages/torch/utils/data/distributed.py", line 74, in __init__
raise ValueError(
ValueError: Invalid rank 2, rank should be in the interval [0, 1]
```

### Environment

```

Current environment
* CUDA:
- GPU:
- NVIDIA A40
- NVIDIA A40
- available: True
- version: 12.1
* Lightning:
- lightning: 2.2.2
- lightning-cloud: 0.5.57
- lightning-utilities: 0.10.1
- open-clip-torch: 2.16.2
- pytorch-lightning: 2.1.3
- torch: 2.1.2
- torchaudio: 2.0.0
- torchmetrics: 0.11.4
- torchvision: 0.15.0
* System:
- OS: Linux
- architecture:
- 64bit
- ELF
- processor: x86_64
- python: 3.10.10
- release: 5.14.0-162.6.1.el9_1.x86_64
- version: #1 SMP PREEMPT_DYNAMIC Fri Sep 30 07:36:03 EDT 2022

#- How you installed Lightning(`conda`, `pip`, source): pip
#- Running environment of LightningApp (e.g. local, cloud): on-prem cluster managed by LSF.

```

### More info

Looking at `ddp.py`, it currently sets `num_replicas` to be:
```
@property
@override
def distributed_sampler_kwargs(self) -> Dict[str, Any]:
return {"num_replicas": (self.num_nodes * self.num_processes), "rank": self.global_rank}
```

Why not simply set `num_replicas` to be `self.world_size`?

[related](https://lightning.ai/forums/t/why-num-replica-world-size/7305?u=shaibagon)

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 in ddp.py at distributed_sampler_kwargs, then trace fabric.py through setup_dataloaders and _get_distributed_sampler. Reproduce the uneven two-node setup described in the issue and verify that ddp and fsdp use consistent replica and rank values without the DistributedSampler rank error.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.