Lightning-AI / Lightning-AI/pytorch-lightning
Device parsing improvements
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
## Proposed refactoring or deprecation
Introduce a device selection dataclass that holds the device selection in a standardized format. Idea by @ananthsub
### Motivation
We have a `_parse_devices` function used in the Trainer and Lite that returns a tuple of parsed device indices.
https://github.com/PyTorchLightning/pytorch-lightning/blob/9237106451f97393b17009a0ca571b6ff5ba5484/pytorch_lightning/trainer/trainer.py#L1459-L1470
From @ananthsub in https://github.com/PyTorchLightning/pytorch-lightning/pull/10230#discussion_r738806860
> returning a tuple isn't going to scale well with more device types. it's not easy to tell which positional index maps to which device id type. it could be better to introduce a dataclass to represent the schema concretely. that would also naturally allow for extensions like IPUs
### Pitch
```py
@dataclass
class DeviceSelection
devices: List[int] = []
type: DeviceType = CPU
def parse_input(gpus, tpu_cores, ipus, ...)
# validate user inputs
# map various input formats to standardized one in this dataclass
...
return DeviceSelection(devices=..., type=...)
```
The AcceleratorConnector new gets as input the DeviceSelection instance instead of a growing list of arguments. It currently takes: devices, gpus, gpu_ids, tpu_cores, ipus, num_processes
### Additional context
Alternative to #10231
#### If you enjoy Lightning, check out our other projects! ⚡
- [**Metrics**](https://github.com/PyTorchLightning/metrics): Machine learning metrics for distributed, scalable PyTorch applications.
- [**Flash**](https://github.com/PyTorchLightning/lightning-flash): The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, finetuning 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 @akihironitta @rohitgr7 @tchaton @borda
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 the _parse_devices function in pytorch_lightning/trainer/trainer.py and trace how Trainer and Lite pass device arguments into AcceleratorConnector. Review the proposed DeviceSelection dataclass and the current devices, gpus, gpu_ids, tpu_cores, ipus, and num_processes inputs. Done means device selection is represented consistently and the affected integration points and validations work across the supported device types.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 20/100