Lightning-AI / Lightning-AI/pytorch-lightning
Refactor distributed_sampler update in `data_connector.py` for distributed strategy
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
## Proposed refactor
This issue follows @ananthsub 's https://github.com/PyTorchLightning/pytorch-lightning/issues/11756 to move strategy-specific dataloader logic to the stategies.
### Motivation
`_prepare_dataloader()` function in `data_connector.py` handles the logic whether the dataloader needs to update the distributed sampler when user is using distributed strategy (see the following code link)
https://github.com/PyTorchLightning/pytorch-lightning/blob/cc43d07db1ab77385feff04c01f040c5cad805a9/pytorch_lightning/trainer/connectors/data_connector.py#L360-L372
However, modifying sampler in dataloader can be pushed directly to strategy class. Similar to issue #12216 , we can move the logic to `Strategy` base class. For any strategy classes have `is_distributed=True` flag, they can verify if the dataloader need to update the sampler.
### Pitch
We can move the logic to `Strategy` base class. For any strategy classes have `is_distributed=True` flag, they can verify if the dataloader need to update the sampler.
### Additional context
cc: @edward-io @four4fish @ananthsub
______________________________________________________________________
#### If you enjoy Lightning, check out our other projects! ⚡
- [**Metrics**](https://github.com/PyTorchLightning/metrics): Machine learning metrics for distributed, scalable PyTorch applications.
- [**Lite**](https://pytorch-lightning.readthedocs.io/en/latest/starter/lightning_lite.html): enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.
- [**Flash**](https://github.com/PyTorchLightning/lightning-flash): The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.
- [**Bolts**](https://github.com/PyTorchLightning/lightning-bolts): Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.
- [**Lightning Transformers**](https://github.com/PyTorchLightning/lightning-transformers): Flexible interface for high-performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.
cc @justusschock @awaelchli @rohitgr7 @ninginthecloud @otaj @tchaton @akihironitta @carmocca
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 with _prepare_dataloader() in pytorch_lightning/trainer/connectors/data_connector.py, especially the linked sampler-update section, then read the Strategy base class and issue #12216 for the related direction. The work is done when the distributed-sampler decision is handled by the appropriate strategy classes while preserving the existing dataloader behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100