Lightning-AI / Lightning-AI/pytorch-lightning

Device parsing improvements

Open
#10,265 4 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

accelerator design refactor
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

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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.