Lightning-AI / Lightning-AI/pytorch-lightning
Returning num_replicas=world_size when using distributed sampler in ddp
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
The default `LightningEnvironment` assumes that every node in a multi-node environment has equal number of GPUs i.e. each node assumes that the [world size is equal to number of nodes multiplied by the number of (active) devices on that node](https://github.com/Lightning-AI/pytorch-lightning/blob/06ea3a05716a6d1f4a96cfb25021accdd18d8146/src/lightning/pytorch/strategies/ddp.py#L211).
However, implementing one's own environment can bypass this limitation (example attached below). While the processes get registered successfully, the attribute `num_replicas` that is provided to the `DistributedSampler` class is still [initialized independently](https://github.com/Lightning-AI/pytorch-lightning/blob/06ea3a05716a6d1f4a96cfb25021accdd18d8146/src/lightning/pytorch/strategies/ddp.py#L137) of the environment, which leads to an error of having ranks outside the scope of the world size.
Fix: Use `num_replicas=self.world_size()` instead of estimating the world size again.
### What version are you seeing the problem on?
v2.2
### How to reproduce the bug
Config and custom Environment:
```python
# Code that sets some config parameters as follows
config.nodes = [
("NODE_NAME", ["NODE_IP_ADDRESS", "NUM_GPUS: int"]), # Index 0 is master node
...
]
config.num_nodes = len(config.nodes)
# Set some global variables
MASTER_PORT = 10051 # Set port here
MASTER_ADDR = config.nodes[0][1][0] # Set address here
WORLD_SIZE = sum([node_info[1] for _, node_info in config.nodes]) # Set world size here
# Set config devices, NODE_RANK, and global rank starting point
NODE_RANK = ...
GLOBAL_RANK_OFFSET = 0
for i, (node, node_info) in enumerate(config.nodes):
if node == socket.gethostname():
config.devices = node_info[1]
NODE_RANK = i
break
GLOBAL_RANK_OFFSET += node_info[1]
# Set environment variables
os.environ["MASTER_ADDR"] = str(MASTER_ADDR)
os.environ["MASTER_PORT"] = str(MASTER_PORT)
os.environ["WORLD_SIZE"] = str(WORLD_SIZE)
os.environ["NODE_RANK"] = str(NODE_RANK)
class MyClusterEnvironment(LightningEnvironment):
def set_world_size(self, size: int):
# Here, size = num_nodes * len(devices) which does not work for heterogenous clusters
self._world_size = WORLD_SIZE
def set_global_rank(self, rank: int):
# Here, global_rank = node_rank * len(devices) + local_rank which does not work for heterogenous clusters
global_rank = GLOBAL_RANK_OFFSET + self.local_rank()
self._global_rank = global_rank
config.cluster_environment = MyClusterEnvironment()
```
Trainer:
```python
trainer = L.Trainer(
num_nodes=config.num_nodes,
devices=config.devices,
plugins=[config.cluster_environment],
... # Other arguments
)
```
Run on:
- Node 0: 4xA100 80GB PCIe GPUs
- Node 1: 2xA100 80GB PCIe GPUs
### Error messages and logs
Error on node 1:
```
╭────────────────────────────────────────────── Traceback (most recent call last) ───────────────────────────────────────────────╮
│ /home/users/arjun.agarwal/projects/mock_training/distributed.py:99 in │
│ │
│ 96 │ │ plugins=[config.cluster_environment], │
│ 97 │ ) │
│ 98 │ │
│ ❱ 99 │ trainer.fit(model, dm) │
│ 100 │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py:544 in fit │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/lightning/pytorch/trainer/call.py:43 in │
│ _call_and_handle_interrupt │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/lightning/pytorch/strategies/launchers/subprocess_script.py:1 │
│ 05 in launch │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py:580 in _fit_impl │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py:987 in _run │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py:1031 in _run_stage │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/lightning/pytorch/trainer/trainer.py:1060 in │
│ _run_sanity_check │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/lightning/pytorch/loops/utilities.py:182 in _decorator │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/lightning/pytorch/loops/evaluation_loop.py:110 in run │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/lightning/pytorch/loops/evaluation_loop.py:180 in setup_data │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:501 in │
│ _process_dataloader │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:206 in │
│ _prepare_dataloader │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:217 in │
│ _resolve_sampler │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:258 in │
│ _get_distributed_sampler │
│ │
│ /home/users/arjun.agarwal/miniconda3/lib/python3.9/site-packages/torch/utils/data/distributed.py:74 in __init__ │
╰────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
ValueError: Invalid rank 4, rank should be in the interval [0, 3]
```
`num_replicas` gets set to 4 here as `num_nodes=2` and `num_processes=2`. However world size is 6 as defined in the environment.
### Environment
Current environment
```
* Lightning:
- efficientnet-pytorch: 0.7.1
- lightning: 2.2.5
- lightning-cloud: 0.5.61
- lightning-utilities: 0.10.0
- pytorch-lightning: 2.1.2
- pytorchvideo: 0.1.5
- torch: 2.2.2
- torchaudio: 2.2.2
- torchmetrics: 1.2.1
- torchsummary: 1.5.1
- torchvision: 0.17.2
```
### More info
The issue can be fixed by replacing [ddp.py:L137](https://github.com/Lightning-AI/pytorch-lightning/blob/06ea3a05716a6d1f4a96cfb25021accdd18d8146/src/lightning/pytorch/strategies/ddp.py#L137)
```python
@property
@override
def distributed_sampler_kwargs(self) -> Dict[str, Any]:
return {"num_replicas": (self.num_nodes * self.num_processes), "rank": self.global_rank}
```
with
```python
@property
@override
def distributed_sampler_kwargs(self) -> Dict[str, Any]:
return {"num_replicas": self.world_size, "rank": self.global_rank}
```
cc @borda @awaelchli @justusschock
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 in src/lightning/pytorch/strategies/ddp.py at distributed_sampler_kwargs, especially the current num_replicas calculation and the environment's world-size access. Reproduce the heterogeneous 4-GPU/2-GPU setup described in the issue and verify that DistributedSampler accepts all global ranks. Done means the sampler uses the configured world size without producing the reported invalid-rank error.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100