Lightning-AI / Lightning-AI/pytorch-lightning

Add device information to the accelerator config message

Open
#17,355 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

accelerator feature help wanted trainer
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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.