Lightning-AI / Lightning-AI/pytorch-lightning
Add device information to the accelerator config message
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Description & Motivation
Revamp
```bash
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
```
To
```bash
GPU available: M1, using 1 devices
TPU available: v4-8, using 0 devices
IPU available: False, using 0 devices
HPU available: False, using 0 devices
```
The relevant code is: https://github.com/Lightning-AI/lightning/blob/f14ee9edbc8269054e12daf30b8681d530e73369/src/lightning/pytorch/trainer/setup.py#L145-L171
### Pitch
If the accelerator is available, `True` changes to the actual name of the accelerator used.
If it's unavailable, we still show `False`.
For GPUs, the `cuda|mps` field is gone, as it should be clear from the device.
I also propose that the GPU field shows the number of devices, instead of a used boolean.
We can get this info via
```python
# CUDA
torch.cuda.get_device_name()
# TPU
from torch_xla.experimental import tpu
import torch_xla.core.xla_env_vars as xenv
# note: this needs a try-except as this will send a request
tpu.get_tpu_env()[xenv.ACCELERATOR_TYPE]
```
For MPS, HPU, IPU we would need to find out if we can get this information. In the meantime, we can still fallback to "True" for them.
This could be done by introducing an `Accelerator.device_name(device)` staticmethod
### Alternatives
One caveat is that this might be misleading with heterogeneous devices, as only rank zero prints this information.
### Additional context
_No response_
cc @borda @justusschock @awaelchli
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 linked section of src/lightning/pytorch/trainer/setup.py and inspect the proposed Accelerator.device_name(device) entry point. Confirm how CUDA, TPU, MPS, HPU, and IPU names and device counts should be obtained, including fallback behavior, then ensure the logged messages match the requested examples.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- ai, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100