Bug of SingleDeviceStrategy: incoherent device between accelerator and strategy when accelerator="auto"

Bug description

Hello lighting team,

I observed a strange behaviour of SingleDeviceStrategy: when using Trainer(accelerator="auto", strategy="single_device"), the strategy's accelerator is not correctly detected.

For example, the following code snippet:

    trainer = Trainer(accelerator="auto", strategy="single_device")
    print("accelerator:", trainer.accelerator)
    print("strategy's root device:", trainer.strategy.root_device)

When running on a machine with GPU, I got:

accelerator: <lightning.pytorch.accelerators.cuda.CUDAAccelerator object at 0x7faefe6c4890>
strategy's root device: cpu

There's a mismatch between the accelerator and strategy's device, and will raise this error when running inference:

lightning.fabric.utilities.exceptions.MisconfigurationException: Device should be GPU, got cpu instead

What version are you seeing the problem on?

v2.0

How to reproduce the bug

from lightning import LightningModule
from lightning.pytorch import Trainer
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset

class SimpleDataset(Dataset):
    def __len__(self):
        return 1000

    def __getitem__(self, idx):
        return torch.randn((1, 28, 28))

class SimpleModel(LightningModule):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.layer1 = nn.Linear(28 * 28, 512)
        self.layer2 = nn.Linear(512, 512)
        self.layer3 = nn.Linear(512, 512)
        self.layer4 = nn.Linear(512, 512)
        self.layer5 = nn.Linear(512, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        x = F.relu(self.layer3(x))
        x = F.relu(self.layer4(x))
        x = self.layer5(x)
        return x


    def predict_step(self, batch, batch_idx, dataloader_idx=None):
        return self(batch)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters())



if __name__ == '__main__':
    trainer = Trainer(accelerator="auto", strategy="single_device")
    print("accelerator", trainer.accelerator)
    print("strategy's root device", trainer.strategy.root_device)

    dataset = SimpleDataset()
    dataloader = DataLoader(dataset, batch_size=32)
    model = SimpleModel()

    trainer = Trainer(strategy="single_device")
    trainer.predict(model, dataloader)

Error messages and logs

# Error messages and logs here please
zekun_zhang_gleamer_ai@gpu-2:~/ai-gleamer$ poetry run python minimal_reproducible_examples/device_conflict_for_single_device_strategy.py
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
accelerator: <lightning.pytorch.accelerators.cuda.CUDAAccelerator object at 0x7faefe6c4890>
strategy's root device: cpu
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
Traceback (most recent call last):
  File "/home/zekun_zhang_gleamer_ai/ai-gleamer/minimal_reproducible_examples/device_conflict_for_single_device_strategy.py", line 52, in <module>
    trainer.predict(model, dataloader)
  File "/home/zekun_zhang_gleamer_ai/.cache/pypoetry/virtualenvs/glmlib-BFgLEs1Y-py3.11/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 805, in predict
    return call._call_and_handle_interrupt(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zekun_zhang_gleamer_ai/.cache/pypoetry/virtualenvs/glmlib-BFgLEs1Y-py3.11/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py", line 44, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zekun_zhang_gleamer_ai/.cache/pypoetry/virtualenvs/glmlib-BFgLEs1Y-py3.11/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 847, in _predict_impl
    results = self._run(model, ckpt_path=ckpt_path)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zekun_zhang_gleamer_ai/.cache/pypoetry/virtualenvs/glmlib-BFgLEs1Y-py3.11/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 893, in _run
    self.strategy.setup_environment()
  File "/home/zekun_zhang_gleamer_ai/.cache/pypoetry/virtualenvs/glmlib-BFgLEs1Y-py3.11/lib/python3.11/site-packages/lightning/pytorch/strategies/strategy.py", line 127, in setup_environment
    self.accelerator.setup_device(self.root_device)
  File "/home/zekun_zhang_gleamer_ai/.cache/pypoetry/virtualenvs/glmlib-BFgLEs1Y-py3.11/lib/python3.11/site-packages/lightning/pytorch/accelerators/cuda.py", line 42, in setup_device
    raise MisconfigurationException(f"Device should be GPU, got {device} instead")
lightning.fabric.utilities.exceptions.MisconfigurationException: Device should be GPU, got cpu instead

Environment

More info

Sorry to bother you again @awaelchli, but you've given great help to me in another issue. I'am grateful for any ideas and suggestions ! 😉

cc @tchaton

Reproduce the mismatch with Trainer(accelerator="auto", strategy="single_device") using the provided minimal example. Start from lightning/pytorch/strategies/strategy.py at setup_environment and lightning/pytorch/accelerators/cuda.py, then verify that the accelerator and strategy root_device agree and trainer.predict completes without the reported exception.