Repository: Lightning-AI/pytorch-lightning
Pull request: Add support for pluggable Accelerators
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

Supports #10687

```python
trainer = Trainer(accelerator=NewSOTAAccelerator(), devices=4)
```

Also, fixes #12044 and #10592

### Does your PR introduce any breaking changes? If yes, please list them.

<!-- FILL IN or None -->

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [ ] Did you make sure to **update the documentation** with your changes? (if necessary)
- [x] Did you write any **new necessary tests**? (not for typos and docs)
- [ ] Did you verify new and **existing tests pass** locally with your changes?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [x] Did you **update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)**? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?

Make sure you had fun coding 🙃


cc @borda @akihironitta @rohitgr7 @justusschock

Problem statement:
Import error when trying to run TPU accelerator on CPU
## 🐛 Bug

<!-- A clear and concise description of the bug. -->

### To Reproduce

```py

import os

import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("valid_loss", loss)

    def test_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("test_loss", loss)

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


def run():
    train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    val_data = DataLoader(RandomDataset(32, 64), batch_size=2)
    test_data = DataLoader(RandomDataset(32, 64), batch_size=2)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        limit_test_batches=1,
        num_sanity_val_steps=0,
        max_epochs=1,
        enable_model_summary=False,
        accelerator="tpu",
        devices=1,
    )
    trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)
    trainer.test(model, dataloaders=test_data)


if __name__ == "__main__":
    run()

```

Produces

```
Traceback (most recent call last):
  File "/Users/adrian/repositories/pytorch-lightning/pl_examples/bug_report/bug_report_model.py", line 68, in <module>
    run()
  File "/Users/adrian/repositories/pytorch-lightning/pl_examples/bug_report/bug_report_model.py", line 52, in run
    trainer = Trainer(
  File "/Users/adrian/repositories/pytorch-lightning/pytorch_lightning/utilities/argparse.py", line 336, in insert_env_defaults
    return fn(self, **kwargs)
  File "/Users/adrian/repositories/pytorch-lightning/pytorch_lightning/trainer/trainer.py", line 477, in __init__
    self._accelerator_connector = AcceleratorConnector(
  File "/Users/adrian/repositories/pytorch-lightning/pytorch_lightning/trainer/connectors/accelerator_connector.py", line 194, in __init__
    self._strategy_flag = self._choose_strategy()
  File "/Users/adrian/repositories/pytorch-lightning/pytorch_lightning/trainer/connectors/accelerator_connector.py", line 531, in _choose_strategy
    return SingleTPUStrategy(device=self._parallel_devices[0])  # type: ignore
  File "/Users/adrian/repositories/pytorch-lightning/pytorch_lightning/strategies/single_tpu.py", line 44, in __init__
    device=xm.xla_device(device),
NameError: name 'xm' is not defined
```

### Expected behavior

Should error with
```

File "/Users/adrian/repositories/pytorch-lightning/pytorch_lightning/trainer/connectors/accelerator_connector.py", line 218, in select_accelerator_type
    raise MisconfigurationException(f"You passed `accelerator='tpu'`, but {msg}.")
pytorch_lightning.utilities.exceptions.MisconfigurationException: You passed `accelerator='tpu'`, but TPUs are not available.

```
like on 1.5.10

### Environment

Latest master.
Ask if you need more.

cc @kaushikb11 @rohitgr7

Hints from issue discussion:
Likely a result of #11448 cc @four4fish 
It would be fixed by the updates added in #12030, and I have moved the parsing logic to the accelerators.
But would do a quick PR for a fix for it
Yeah, this is the follow up item "Enable accelerator.is_available() check" in #11449 
The proper way is to call `accelerator.is_available()` in `_init_acceleartor()`, left it as a follow up because it cause a lot of GPU tests fails. The previous accl_con logic, device availability check doesn't apply to GPU. Now calling `self.accelerator.is_available()` will apply device check to GPU as well, and a lot of tests need adding mocks. To avoid massive tests change in one PR, I left this as follow up.
Is this urgent? Do you prefer a quick fix or waiting for  "Enable accelerator.is_available() check" to fix this?
@awaelchli @kaushikb11

Commit message: Update pytorch_lightning/trainer/connectors/accelerator_connector.py

Co-authored-by: Rohit Gupta <rohitgr1998@gmail.com>

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2a2abb489d6cb..7b14f895c8108 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -131,6 +131,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added utility functions for moving optimizers to devices ([#11758](https://github.com/PyTorchLightning/pytorch-lightning/pull/11758))
 
 
+- Added support for pluggable Accelerators ([#12030](https://github.com/PyTorchLightning/pytorch-lightning/pull/12030))
+
 
 ### Changed
 
diff --git a/docs/source/extensions/accelerator.rst b/docs/source/extensions/accelerator.rst
index 03c547f49a0f8..762f2b6e57a90 100644
--- a/docs/source/extensions/accelerator.rst
+++ b/docs/source/extensions/accelerator.rst
@@ -20,6 +20,7 @@ Each Accelerator gets two plugins upon initialization:
 One to handle differences from the training routine and one to handle different precisions.
 
 .. testcode::
+    :skipif: torch.cuda.device_count() < 2
 
     from pytorch_lightning import Trainer
     from pytorch_lightning.accelerators import GPUAccelerator
@@ -28,8 +29,8 @@ One to handle differences from the training routine and one to handle different
 
     accelerator = GPUAccelerator()
     precision_plugin = NativeMixedPrecisionPlugin(precision=16, device="cuda")
-    training_type_plugin = DDPStrategy(accelerator=accelerator, precision_plugin=precision_plugin)
-    trainer = Trainer(strategy=training_type_plugin)
+    training_strategy = DDPStrategy(accelerator=accelerator, precision_plugin=precision_plugin)
+    trainer = Trainer(strategy=training_strategy, devices=2)
 
 
 We expose Accelerators and Plugins mainly for expert users who want to extend Lightning to work with new
diff --git a/pytorch_lightning/accelerators/accelerator.py b/pytorch_lightning/accelerators/accelerator.py
index 06f82fb8d4b96..cbd0e2309e311 100644
--- a/pytorch_lightning/accelerators/accelerator.py
+++ b/pytorch_lightning/accelerators/accelerator.py
@@ -55,6 +55,16 @@ def get_device_stats(self, device: Union[str, torch.device]) -> Dict[str, Any]:
         """
         raise NotImplementedError
 
+    @staticmethod
+    @abstractmethod
+    def parse_devices(devices: Any) -> Any:
+        """Accelerator device parsing logic."""
+
+    @staticmethod
+    @abstractmethod
+    def get_parallel_devices(devices: Any) -> Any:
+        """Gets parallel devices for the Accelerator."""
+
     @staticmethod
     @abstractmethod
     def auto_device_count() -> int:
diff --git a/pytorch_lightning/accelerators/cpu.py b/pytorch_lightning/accelerators/cpu.py
index 2fbe3bf18b079..7ab1ebf67c0c2 100644
--- a/pytorch_lightning/accelerators/cpu.py
+++ b/pytorch_lightning/accelerators/cpu.py
@@ -11,14 +11,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from __future__ import annotations
-
-from typing import Any
+from typing import Any, List, Union
 
 import torch
 
 from pytorch_lightning.accelerators.accelerator import Accelerator
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.rank_zero import rank_zero_warn
 from pytorch_lightning.utilities.types import _DEVICE
 
 
@@ -39,6 +38,19 @@ def get_device_stats(self, device: _DEVICE) -> dict[str, Any]:
         """CPU device stats aren't supported yet."""
         return {}
 
+    @staticmethod
+    def parse_devices(devices: Union[int, str, List[int]]) -> Union[int, str, List[int]]:
+        """Accelerator device parsing logic."""
+        return devices
+
+    @staticmethod
+    def get_parallel_devices(devices: Union[int, str, List[int]]) -> List[torch.device]:
+        """Gets parallel devices for the Accelerator."""
+        if isinstance(devices, int):
+            return [torch.device("cpu")] * devices
+        rank_zero_warn(f"The flag `devices` must be an int with `accelerator='cpu'`, got `devices={devices}` instead.")
+        return []
+
     @staticmethod
     def auto_device_count() -> int:
         """Get the devices when set to auto."""
diff --git a/pytorch_lightning/accelerators/gpu.py b/pytorch_lightning/accelerators/gpu.py
index aa8b0d56dbf63..6b767b2d29e23 100644
--- a/pytorch_lightning/accelerators/gpu.py
+++ b/pytorch_lightning/accelerators/gpu.py
@@ -11,18 +11,17 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from __future__ import annotations
-
 import logging
 import os
 import shutil
 import subprocess
-from typing import Any
+from typing import Any, List, Optional, Union
 
 import torch
 
 import pytorch_lightning as pl
 from pytorch_lightning.accelerators.accelerator import Accelerator
+from pytorch_lightning.utilities import device_parser
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.imports import _TORCH_GREATER_EQUAL_1_8
 from pytorch_lightning.utilities.types import _DEVICE
@@ -75,6 +74,16 @@ def get_device_stats(self, device: _DEVICE) -> dict[str, Any]:
             return torch.cuda.memory_stats(device)
         return get_nvidia_gpu_stats(device)
 
+    @staticmethod
+    def parse_devices(devices: Union[int, str, List[int]]) -> Optional[List[int]]:
+        """Accelerator device parsing logic."""
+        return device_parser.parse_gpu_ids(devices)
+
+    @staticmethod
+    def get_parallel_devices(devices: List[int]) -> List[torch.device]:
+        """Gets parallel devices for the Accelerator."""
+        return [torch.device("cuda", i) for i in devices]
+
     @staticmethod
     def auto_device_count() -> int:
         """Get the devices when set to auto."""
diff --git a/pytorch_lightning/accelerators/ipu.py b/pytorch_lightning/accelerators/ipu.py
index 6928546cf8c50..2ac1c794610d8 100644
--- a/pytorch_lightning/accelerators/ipu.py
+++ b/pytorch_lightning/accelerators/ipu.py
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, Union
+from typing import Any, Dict, List, Union
 
 import torch
 
@@ -26,6 +26,16 @@ def get_device_stats(self, device: Union[str, torch.device]) -> Dict[str, Any]:
         """IPU device stats aren't supported yet."""
         return {}
 
+    @staticmethod
+    def parse_devices(devices: int) -> int:
+        """Accelerator device parsing logic."""
+        return devices
+
+    @staticmethod
+    def get_parallel_devices(devices: int) -> List[int]:
+        """Gets parallel devices for the Accelerator."""
+        return list(range(devices))
+
     @staticmethod
     def auto_device_count() -> int:
         """Get the devices when set to auto."""
diff --git a/pytorch_lightning/accelerators/tpu.py b/pytorch_lightning/accelerators/tpu.py
index f1f598c3f1b3c..cd84cccd8b493 100644
--- a/pytorch_lightning/accelerators/tpu.py
+++ b/pytorch_lightning/accelerators/tpu.py
@@ -11,11 +11,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Dict, Union
+from typing import Any, Dict, List, Optional, Union
 
 import torch
 
 from pytorch_lightning.accelerators.accelerator import Accelerator
+from pytorch_lightning.utilities import device_parser
 from pytorch_lightning.utilities.imports import _TPU_AVAILABLE, _XLA_AVAILABLE
 
 if _XLA_AVAILABLE:
@@ -43,6 +44,18 @@ def get_device_stats(self, device: Union[str, torch.device]) -> Dict[str, Any]:
         }
         return device_stats
 
+    @staticmethod
+    def parse_devices(devices: Union[int, str, List[int]]) -> Optional[Union[int, List[int]]]:
+        """Accelerator device parsing logic."""
+        return device_parser.parse_tpu_cores(devices)
+
+    @staticmethod
+    def get_parallel_devices(devices: Union[int, List[int]]) -> List[int]:
+        """Gets parallel devices for the Accelerator."""
+        if isinstance(devices, int):
+            return list(range(devices))
+        return devices
+
     @staticmethod
     def auto_device_count() -> int:
         """Get the devices when set to auto."""
diff --git a/pytorch_lightning/trainer/connectors/accelerator_connector.py b/pytorch_lightning/trainer/connectors/accelerator_connector.py
index f2d27a249f6f2..5a49bec28e453 100644
--- a/pytorch_lightning/trainer/connectors/accelerator_connector.py
+++ b/pytorch_lightning/trainer/connectors/accelerator_connector.py
@@ -163,6 +163,7 @@ def __init__(
         self._precision_flag: Optional[Union[int, str]] = None
         self._precision_plugin_flag: Optional[PrecisionPlugin] = None
         self._cluster_environment_flag: Optional[Union[ClusterEnvironment, str]] = None
+        self._parallel_devices: List[Union[int, torch.device]] = []
         self.checkpoint_io: Optional[CheckpointIO] = None
         self._amp_type_flag: Optional[LightningEnum] = None
         self._amp_level_flag: Optional[str] = amp_level
@@ -356,6 +357,7 @@ def _check_config_and_set_final_flags(
                         self._accelerator_flag = "cpu"
                     if self._strategy_flag.parallel_devices[0].type == "cuda":
                         self._accelerator_flag = "gpu"
+                    self._parallel_devices = self._strategy_flag.parallel_devices
 
         amp_type = amp_type if isinstance(amp_type, str) else None
         self._amp_type_flag = AMPType.from_str(amp_type)
@@ -382,7 +384,7 @@ def _check_device_config_and_set_final_flags(
             devices, num_processes, gpus, ipus, tpu_cores
         )
 
-        if self._devices_flag in ([], 0, "0", "0,"):
+        if self._devices_flag in ([], 0, "0"):
             rank_zero_warn(f"You passed `devices={devices}`, switching to `cpu` accelerator")
             self._accelerator_flag = "cpu"
 
@@ -403,10 +405,8 @@ def _map_deprecated_devices_specfic_info_to_accelerator_and_device_flag(
         """Sets the `devices_flag` and `accelerator_flag` based on num_processes, gpus, ipus, tpu_cores."""
         self._gpus: Optional[Union[List[int], str, int]] = gpus
         self._tpu_cores: Optional[Union[List[int], str, int]] = tpu_cores
-        gpus = device_parser.parse_gpu_ids(gpus)
-        tpu_cores = device_parser.parse_tpu_cores(tpu_cores)
         deprecated_devices_specific_flag = num_processes or gpus or ipus or tpu_cores
-        if deprecated_devices_specific_flag and deprecated_devices_specific_flag not in (0, "0"):
+        if deprecated_devices_specific_flag and deprecated_devices_specific_flag not in ([], 0, "0"):
             if devices:
                 # TODO: @awaelchli improve error message
                 rank_zero_warn(
@@ -451,51 +451,33 @@ def _choose_accelerator(self) -> str:
 
     def _set_parallel_devices_and_init_accelerator(self) -> None:
         # TODO add device availability check
-        self._parallel_devices: List[Union[int, torch.device]] = []
-
         if isinstance(self._accelerator_flag, Accelerator):
             self.accelerator: Accelerator = self._accelerator_flag
-        elif self._accelerator_flag == "tpu":
-            self.accelerator = TPUAccelerator()
-            self._set_devices_flag_if_auto_passed()
-            if isinstance(self._devices_flag, int):
-                self._parallel_devices = list(range(self._devices_flag))
-            else:
-                self._parallel_devices = self._devices_flag  # type: ignore[assignment]
-
-        elif self._accelerator_flag == "ipu":
-            self.accelerator = IPUAccelerator()
-            self._set_devices_flag_if_auto_passed()
-            if isinstance(self._devices_flag, int):
-                self._parallel_devices = list(range(self._devices_flag))
-
-        elif self._accelerator_flag == "gpu":
-            self.accelerator = GPUAccelerator()
-            self._set_devices_flag_if_auto_passed()
-            if isinstance(self._devices_flag, int) or isinstance(self._devices_flag, str):
-                self._devices_flag = int(self._devices_flag)
-                self._parallel_devices = (
-                    [torch.device("cuda", i) for i in device_parser.parse_gpu_ids(self._devices_flag)]  # type: ignore
-                    if self._devices_flag != 0
-                    else []
+        else:
+            ACCELERATORS = {
+                "cpu": CPUAccelerator,
+                "gpu": GPUAccelerator,
+                "tpu": TPUAccelerator,
+                "ipu": IPUAccelerator,
+            }
+            self._accelerator_flag = self._accelerator_flag.lower()  # type: ignore
+            if self._accelerator_flag not in ACCELERATORS:
+                raise MisconfigurationException(
+                    "When passing string value for the `accelerator` argument of `Trainer`,"
+                    f" it can only be one of {list(ACCELERATORS)}."
                 )
-            else:
-                self._parallel_devices = [torch.device("cuda", i) for i in self._devices_flag]  # type: ignore
+            accelerator_class = ACCELERATORS[self._accelerator_flag]
+            self.accelerator = accelerator_class()  # type: ignore
 
-        elif self._accelerator_flag == "cpu":
-            self.accelerator = CPUAccelerator()
-            self._set_devices_flag_if_auto_passed()
-            if isinstance(self._devices_flag, int):
-                self._parallel_devices = [torch.device("cpu")] * self._devices_flag
-            else:
-                rank_zero_warn(
-                    "The flag `devices` must be an int with `accelerator='cpu'`,"
-                    f" got `devices={self._devices_flag}` instead."
-                )
+        self._set_devices_flag_if_auto_passed()
 
         self._gpus = self._devices_flag if not self._gpus else self._gpus
         self._tpu_cores = self._devices_flag if not self._tpu_cores else self._tpu_cores
 
+        self._devices_flag = self.accelerator.parse_devices(self._devices_flag)
+        if not self._parallel_devices:
+            self._parallel_devices = self.accelerator.get_parallel_devices(self._devices_flag)
+
     def _set_devices_flag_if_auto_passed(self) -> None:
         if self._devices_flag == "auto" or not self._devices_flag:
             self._devices_flag = self.accelerator.auto_device_count()
diff --git a/tests/accelerators/test_accelerator_connector.py b/tests/accelerators/test_accelerator_connector.py
index 0e13b4af0f8d2..ad065b4c0c9d1 100644
--- a/tests/accelerators/test_accelerator_connector.py
+++ b/tests/accelerators/test_accelerator_connector.py
@@ -341,6 +341,14 @@ def creates_processes_externally(self) -> bool:
 @mock.patch("pytorch_lightning.strategies.DDPStrategy.setup_distributed", autospec=True)
 def test_custom_accelerator(device_count_mock, setup_distributed_mock):
     class Accel(Accelerator):
+        @staticmethod
+        def parse_devices(devices):
+            return devices
+
+        @staticmethod
+        def get_parallel_devices(devices):
+            return [torch.device("cpu")] * devices
+
         @staticmethod
         def auto_device_count() -> int:
             return 1
@@ -413,10 +421,17 @@ def test_ipython_incompatible_backend_error(_, monkeypatch):
         Trainer(strategy="dp")
 
 
-@pytest.mark.parametrize("trainer_kwargs", [{}, dict(strategy="dp", accelerator="gpu"), dict(accelerator="tpu")])
-def test_ipython_compatible_backend(trainer_kwargs, monkeypatch):
+@mock.patch("torch.cuda.device_count", return_value=2)
+def test_ipython_compatible_dp_strategy_gpu(_, monkeypatch):
     monkeypatch.setattr(pytorch_lightning.utilities, "_IS_INTERACTIVE", True)
-    trainer = Trainer(**trainer_kwargs)
+    trainer = Trainer(strategy="dp", accelerator="gpu")
+    assert trainer.strategy.launcher is None or trainer.strategy.launcher.is_interactive_compatible
+
+
+@mock.patch("pytorch_lightning.accelerators.tpu.TPUAccelerator.parse_devices", return_value=8)
+def test_ipython_compatible_strategy_tpu(_, monkeypatch):
+    monkeypatch.setattr(pytorch_lightning.utilities, "_IS_INTERACTIVE", True)
+    trainer = Trainer(accelerator="tpu")
     assert trainer.strategy.launcher is None or trainer.strategy.launcher.is_interactive_compatible
 
 
@@ -883,12 +898,9 @@ def test_strategy_choice_ddp_cpu_slurm(device_count_mock, setup_distributed_mock
     assert trainer.strategy.local_rank == 0
 
 
+@mock.patch("pytorch_lightning.accelerators.tpu.TPUAccelerator.parse_devices", return_value=8)
 def test_unsupported_tpu_choice(monkeypatch):
-    import pytorch_lightning.utilities.imports as imports
-    from pytorch_lightning.trainer.connectors.accelerator_connector import AcceleratorConnector
 
-    monkeypatch.setattr(imports, "_XLA_AVAILABLE", True)
-    monkeypatch.setattr(AcceleratorConnector, "has_tpu", True)
     with pytest.raises(MisconfigurationException, match=r"accelerator='tpu', precision=64\)` is not implemented"):
         Trainer(accelerator="tpu", precision=64)
 
diff --git a/tests/accelerators/test_common.py b/tests/accelerators/test_common.py
index 4190553d90115..d2771b15fe8cf 100644
--- a/tests/accelerators/test_common.py
+++ b/tests/accelerators/test_common.py
@@ -13,7 +13,12 @@
 # limitations under the License.
 from unittest import mock
 
-from pytorch_lightning.accelerators import CPUAccelerator, GPUAccelerator, IPUAccelerator, TPUAccelerator
+import torch
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.accelerators import Accelerator, CPUAccelerator, GPUAccelerator, IPUAccelerator, TPUAccelerator
+from pytorch_lightning.strategies import DDPStrategy
+from tests.helpers.boring_model import BoringModel
 
 
 @mock.patch("torch.cuda.device_count", return_value=2)
@@ -22,3 +27,50 @@ def test_auto_device_count(device_count_mock):
     assert GPUAccelerator.auto_device_count() == 2
     assert TPUAccelerator.auto_device_count() == 8
     assert IPUAccelerator.auto_device_count() == 4
+
+
+def test_pluggable_accelerator(tmpdir):
+    class TestAccelerator(Accelerator):
+        @staticmethod
+        def parse_devices(devices):
+            return devices
+
+        @staticmethod
+        def get_parallel_devices(devices):
+            return [torch.device("cpu")] * devices
+
+        @staticmethod
+        def auto_device_count():
+            return 1
+
+        @staticmethod
+        def is_available():
+            return True
+
+    model = BoringModel()
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        limit_train_batches=2,
+        limit_val_batches=2,
+        max_epochs=1,
+        accelerator=TestAccelerator(),
+        devices=2,
+        strategy="ddp",
+    )
+    trainer.fit(model)
+
+    assert isinstance(trainer.accelerator, TestAccelerator)
+    assert trainer._accelerator_connector.parallel_devices == [torch.device("cpu")] * 2
+
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        limit_train_batches=2,
+        limit_val_batches=2,
+        max_epochs=1,
+        strategy=DDPStrategy(TestAccelerator()),
+        devices=2,
+    )
+    trainer.fit(model)
+
+    assert isinstance(trainer.accelerator, TestAccelerator)
+    assert trainer._accelerator_connector.parallel_devices == [torch.device("cpu")] * 2
diff --git a/tests/accelerators/test_tpu.py b/tests/accelerators/test_tpu.py
index 7e1d8c15cf323..f7b031105278e 100644
--- a/tests/accelerators/test_tpu.py
+++ b/tests/accelerators/test_tpu.py
@@ -266,28 +266,30 @@ def forward(self, x):
     assert torch.all(torch.eq(model.net_a.layer.weight, model.net_b.layer.weight))
 
 
+@RunIf(tpu=True)
 def test_tpu_invalid_raises():
     strategy = TPUSpawnStrategy(accelerator=TPUAccelerator(), precision_plugin=PrecisionPlugin())
     with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `TPUPrecisionPlugin"):
-        Trainer(strategy=strategy)
+        Trainer(strategy=strategy, devices=8)
 
     strategy = DDPStrategy(accelerator=TPUAccelerator(), precision_plugin=TPUPrecisionPlugin())
     with pytest.raises(ValueError, match="TPUAccelerator` can only be used with a `SingleTPUStrategy`"):
-        Trainer(strategy=strategy)
+        Trainer(strategy=strategy, devices=8)
 
 
+@RunIf(tpu=True)
 def test_tpu_invalid_raises_set_precision_with_strategy():
     accelerator = TPUAccelerator()
     strategy = TPUSpawnStrategy(accelerator=accelerator, precision_plugin=PrecisionPlugin())
     with pytest.raises(ValueError, match="`TPUAccelerator` can only be used with a `TPUPrecisionPlugin`"):
-        Trainer(strategy=strategy)
+        Trainer(strategy=strategy, devices=8)
 
     accelerator = TPUAccelerator()
     strategy = DDPStrategy(accelerator=accelerator, precision_plugin=TPUPrecisionPlugin())
     with pytest.raises(
         ValueError, match="The `TPUAccelerator` can only be used with a `SingleTPUStrategy` or `TPUSpawnStrategy"
     ):
-        Trainer(strategy=strategy)
+        Trainer(strategy=strategy, devices=8)
 
 
 @RunIf(tpu=True)
diff --git a/tests/conftest.py b/tests/conftest.py
index 8ad7faa3cd769..4440548fa09bc 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -77,6 +77,7 @@ def restore_env_variables():
         "XRT_HOST_WORLD_SIZE",
         "XRT_SHARD_ORDINAL",
         "XRT_SHARD_LOCAL_ORDINAL",
+        "TF_CPP_MIN_LOG_LEVEL",
     }
     leaked_vars.difference_update(allowlist)
     assert not leaked_vars, f"test is leaking environment variable(s): {set(leaked_vars)}"
diff --git a/tests/trainer/test_trainer_cli.py b/tests/trainer/test_trainer_cli.py
index b5713893f769b..8a074e20ee055 100644
--- a/tests/trainer/test_trainer_cli.py
+++ b/tests/trainer/test_trainer_cli.py
@@ -17,11 +17,11 @@
 from unittest import mock
 
 import pytest
+import torch
 
 import tests.helpers.utils as tutils
 from pytorch_lightning import Trainer
 from pytorch_lightning.utilities import argparse
-from tests.helpers.runif import RunIf
 
 
 @mock.patch("argparse.ArgumentParser.parse_args")
@@ -163,11 +163,14 @@ def test_argparse_args_parsing_fast_dev_run(cli_args, expected):
 
 @pytest.mark.parametrize(
     ["cli_args", "expected_parsed", "expected_device_ids"],
-    [("", None, None), ("--accelerator gpu --devices 1", "1", [0]), ("--accelerator gpu --devices 0,", "0,", None)],
+    [("", None, None), ("--accelerator gpu --devices 1", "1", [0]), ("--accelerator gpu --devices 0,", "0,", [0])],
 )
-@RunIf(min_gpus=1)
-def test_argparse_args_parsing_devices(cli_args, expected_parsed, expected_device_ids):
+def test_argparse_args_parsing_devices(cli_args, expected_parsed, expected_device_ids, monkeypatch):
     """Test multi type argument with bool."""
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "device_count", lambda: 1)
+
     cli_args = cli_args.split(" ") if cli_args else []
     with mock.patch("argparse._sys.argv", ["any.py"] + cli_args):
         parser = ArgumentParser(add_help=False)


Code review comments:
diff_hunk: @@ -22,3 +27,50 @@ def test_auto_device_count(device_count_mock):
     assert GPUAccelerator.auto_device_count() == 2
     assert TPUAccelerator.auto_device_count() == 8
     assert IPUAccelerator.auto_device_count() == 4
+
+
+def test_pluggable_accelerator(tmpdir):
+    class TestAccelerator(Accelerator):
+        @staticmethod
+        def parse_devices(devices):
+            return devices
+
+        @staticmethod
+        def get_parallel_devices(devices):
+            return [torch.device("cpu")] * devices
+
+        @staticmethod
+        def auto_device_count():
+            return 1
+
+        @staticmethod
+        def is_available():
+            return True
+
+    model = BoringModel()
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        limit_train_batches=2,
+        limit_val_batches=2,
+        max_epochs=1,
+        accelerator=TestAccelerator(),
+        devices=2,
+        strategy="ddp",
+    )
+    trainer.fit(model)
path: tests/accelerators/test_common.py
text: @user1:
Can we avoid having to call `fit` twice in this test? There's no need to train for these assertions

@author:
Good catch! Updated.

diff_hunk: @@ -39,6 +38,19 @@ def get_device_stats(self, device: _DEVICE) -> dict[str, Any]:
         """CPU device stats aren't supported yet."""
         return {}
 
+    @staticmethod
+    def parse_devices(devices: Union[int, str, List[int]]) -> Union[int, str, List[int]]:
+        """Accelerator device parsing logic."""
+        return devices
+
+    @staticmethod
+    def get_parallel_devices(devices: Union[int, str, List[int]]) -> List[torch.device]:
+        """Gets parallel devices for the Accelerator."""
+        if isinstance(devices, int):
+            return [torch.device("cpu")] * devices
+        rank_zero_warn(f"The flag `devices` must be an int with `accelerator='cpu'`, got `devices={devices}` instead.")
path: pytorch_lightning/accelerators/cpu.py
text: @user1:
```suggestion
        rank_zero_warn(f"The flag `devices` must be an int with `accelerator='cpu'`, got `devices={devices!r}` instead.")
```