Merge different gpu backends with accelerator='gpu'

src/pytorch_lightning/CHANGELOG.md

@@ -110,6 +110,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- `accelerator="gpu"` now automatically selects an available GPU backend (CUDA and MPS currently) ([#13642](https://github.com/Lightning-AI/lightning/pull/13642))
+
+
 - Enable validation during overfitting ([#12527](https://github.com/PyTorchLightning/pytorch-lightning/pull/12527))
 
 
@@ -166,6 +169,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Deprecated
 
+- Deprecated `pytorch_lightning.accelerators.gpu.GPUAccelerator` in favor of `pytorch_lightning.accelerators.cuda.CUDAAccelerator` ([#13636](https://github.com/Lightning-AI/lightning/pull/13636))
+
+
 - Deprecated `pytorch_lightning.loggers.base.LightningLoggerBase` in favor of `pytorch_lightning.loggers.logger.Logger`, and deprecated `pytorch_lightning.loggers.base` in favor of `pytorch_lightning.loggers.logger` ([#120148](https://github.com/PyTorchLightning/pytorch-lightning/pull/12014))
 
 

src/pytorch_lightning/accelerators/cuda.py

@@ -97,12 +97,6 @@ def register_accelerators(cls, accelerator_registry: Dict) -> None:
             cls,
             description=f"{cls.__class__.__name__}",
         )
-        # temporarily enable "gpu" to point to the CUDA Accelerator
-        accelerator_registry.register(
-            "gpu",
-            cls,
-            description=f"{cls.__class__.__name__}",
-        )
 
     def teardown(self) -> None:
         # clean up memory

src/pytorch_lightning/lite/lite.py

@@ -54,7 +54,8 @@ class LightningLite(ABC):
     - Multi-node support.
 
     Args:
-        accelerator: The hardware to run on. Possible choices are: ``"cpu"``, ``"gpu"``, ``"tpu"``, ``"auto"``.
+        accelerator: The hardware to run on. Possible choices are:
+            ``"cpu"``, ``"cuda"``, ``"mps"``, ``"gpu"``, ``"tpu"``, ``"auto"``.
         strategy: Strategy for how to run across multiple devices. Possible choices are:
             ``"dp"``, ``"ddp"``, ``"ddp_spawn"``, ``"deepspeed"``, ``"ddp_sharded"``.
         devices: Number of devices to train on (``int``), which GPUs to train on (``list`` or ``str``), or ``"auto"``.
@@ -436,7 +437,7 @@ def _get_distributed_sampler(dataloader: DataLoader, **kwargs: Any) -> Distribut
         return DistributedSamplerWrapper(dataloader.sampler, **kwargs)
 
     def _check_accelerator_support(self, accelerator: Optional[Union[str, Accelerator]]) -> None:
-        supported = [t.value.lower() for t in self._supported_device_types()] + ["auto"]
+        supported = [t.value.lower() for t in self._supported_device_types()] + ["gpu", "auto"]
         valid = accelerator is None or isinstance(accelerator, Accelerator) or accelerator in supported
         if not valid:
             raise MisconfigurationException(
@@ -457,7 +458,7 @@ def _check_strategy_support(self, strategy: Optional[Union[str, Strategy]]) -> N
     def _supported_device_types() -> Sequence[_AcceleratorType]:
         return (
             _AcceleratorType.CPU,
-            _AcceleratorType.GPU,
+            _AcceleratorType.CUDA,
             _AcceleratorType.TPU,
             _AcceleratorType.MPS,
         )

src/pytorch_lightning/trainer/connectors/accelerator_connector.py

@@ -201,10 +201,14 @@ def __init__(
             devices=devices, num_nodes=num_nodes, num_processes=num_processes, gpus=gpus, ipus=ipus, tpu_cores=tpu_cores
         )
         # 2. Instantiate Accelerator
-        # handle `auto` and `None`
         self._set_accelerator_if_ipu_strategy_is_passed()
+
+        # handle `auto`, `None` and `gpu`
         if self._accelerator_flag == "auto" or self._accelerator_flag is None:
-            self._accelerator_flag = self._choose_accelerator()
+            self._accelerator_flag = self._choose_auto_accelerator()
+        elif self._accelerator_flag == "gpu":
+            self._accelerator_flag = self._choose_gpu_accelerator_backend()
+
         self._set_parallel_devices_and_init_accelerator()
 
         # 3. Instantiate ClusterEnvironment
@@ -280,7 +284,7 @@ def _check_config_and_set_final_flags(
         if (
             accelerator is not None
             and accelerator not in self._accelerator_types
-            and accelerator != "auto"
+            and accelerator not in ("auto", "gpu")
             and not isinstance(accelerator, Accelerator)
         ):
             raise ValueError(
@@ -487,7 +491,7 @@ def _set_accelerator_if_ipu_strategy_is_passed(self) -> None:
         if isinstance(self._strategy_flag, IPUStrategy):
             self._accelerator_flag = "ipu"
 
-    def _choose_accelerator(self) -> str:
+    def _choose_auto_accelerator(self) -> str:
         """Choose the accelerator type (str) based on availability when ``accelerator='auto'``."""
         if self._accelerator_flag == "auto":
             if _TPU_AVAILABLE:
@@ -502,6 +506,15 @@ def _choose_accelerator(self) -> str:
                 return "cuda"
         return "cpu"
 
+    @staticmethod
+    def _choose_gpu_accelerator_backend() -> str:
+        if MPSAccelerator.is_available():
+            return "mps"
+        if CUDAAccelerator.is_available():
+            return "cuda"
+
+        raise MisconfigurationException("No supported gpu backend found!")
+
     def _set_parallel_devices_and_init_accelerator(self) -> None:
         if isinstance(self._accelerator_flag, Accelerator):
             self.accelerator: Accelerator = self._accelerator_flag

src/pytorch_lightning/utilities/enums.py

@@ -244,15 +244,15 @@ class _AcceleratorType(LightningEnum):
     >>> _AcceleratorType.CPU == _AcceleratorType.from_str('cpu')
     True
     >>> # you can match the type with string
-    >>> _AcceleratorType.GPU == 'GPU'
+    >>> _AcceleratorType.CUDA == 'CUDA'
     True
     >>> # which is case invariant
     >>> _AcceleratorType.TPU in ('tpu', 'CPU')
     True
     """
 
     CPU = "CPU"
-    GPU = "GPU"
+    CUDA = "CUDA"
     IPU = "IPU"
     TPU = "TPU"
     HPU = "HPU"

tests/tests_pytorch/accelerators/test_accelerator_connector.py

@@ -278,7 +278,7 @@ def test_accelerator_cpu(_):
         MisconfigurationException,
         match="CUDAAccelerator can not run on your system since the accelerator is not available.",
     ):
-        Trainer(accelerator="gpu")
+        Trainer(accelerator="cuda")
 
     with pytest.deprecated_call(match=r"is deprecated in v1.7 and will be removed"):
         Trainer(accelerator="cpu", gpus=1)
@@ -671,7 +671,7 @@ def test_devices_auto_choice_mps():
 
 @pytest.mark.parametrize(
     ["parallel_devices", "accelerator"],
-    [([torch.device("cpu")], "gpu"), ([torch.device("cuda", i) for i in range(8)], ("tpu"))],
+    [([torch.device("cpu")], "cuda"), ([torch.device("cuda", i) for i in range(8)], ("tpu"))],
 )
 def test_parallel_devices_in_strategy_confilict_with_accelerator(parallel_devices, accelerator):
     with pytest.raises(MisconfigurationException, match=r"parallel_devices set through"):
@@ -746,13 +746,51 @@ def test_plugin_only_one_instance_for_one_type(plugins, expected):
         Trainer(plugins=plugins)
 
 
-@pytest.mark.parametrize("accelerator", ("cpu", "gpu", "tpu", "ipu"))
+@pytest.mark.parametrize("accelerator", ("cpu", "cuda", "mps", "tpu", "ipu"))
 @pytest.mark.parametrize("devices", ("0", 0, []))
 def test_passing_zero_and_empty_list_to_devices_flag(accelerator, devices):
     with pytest.raises(MisconfigurationException, match="value is not a valid input using"):
         Trainer(accelerator=accelerator, devices=devices)
 
 
+@pytest.mark.parametrize(
+    "expected_accelerator_flag,expected_accelerator_class",
+    [
+        pytest.param("cuda", CUDAAccelerator, marks=RunIf(min_cuda_gpus=1)),
+        pytest.param("mps", MPSAccelerator, marks=RunIf(mps=True)),
+    ],
+)
+def test_gpu_accelerator_backend_choice(expected_accelerator_flag, expected_accelerator_class):
+
+    trainer = Trainer(accelerator="gpu")
+    assert trainer._accelerator_connector._accelerator_flag == expected_accelerator_flag
+    assert isinstance(trainer.accelerator, expected_accelerator_class)
+
+
+@mock.patch("pytorch_lightning.utilities.device_parser.num_cuda_devices", return_value=1)
+def test_gpu_accelerator_backend_choice_cuda(_):
+    trainer = Trainer(accelerator="gpu")
+
+    assert trainer._accelerator_connector._accelerator_flag == "cuda"
+    assert isinstance(trainer.accelerator, CUDAAccelerator)
+
+
+@mock.patch("pytorch_lightning.accelerators.mps._MPS_AVAILABLE", return_value=True)
+@mock.patch("torch.device", return_value="mps")  # necessary because torch doesn't allow creation of mps devices
+def test_gpu_accelerator_backend_choice_mps(*_):
+    trainer = Trainer(accelerator="gpu")
+
+    assert trainer._accelerator_connector._accelerator_flag == "mps"
+    assert isinstance(trainer.accelerator, MPSAccelerator)
+
+
+@mock.patch("pytorch_lightning.accelerators.mps.MPSAccelerator.is_available", return_value=False)
+@mock.patch("pytorch_lightning.accelerators.cuda.CUDAAccelerator.is_available", return_value=False)
+def test_gpu_accelerator_misconfiguration_exception(*_):
+    with pytest.raises(MisconfigurationException, match="No supported gpu backend found!"):
+        Trainer(accelerator="gpu")
+
+
 @mock.patch("pytorch_lightning.accelerators.hpu.HPUAccelerator.is_available", return_value=True)
 @mock.patch("pytorch_lightning.strategies.hpu_parallel._HPU_AVAILABLE", return_value=True)
 @mock.patch("pytorch_lightning.plugins.precision.hpu._HPU_AVAILABLE", return_value=True)

tests/tests_pytorch/accelerators/test_accelerator_registry.py

@@ -63,4 +63,4 @@ def is_available():
 
 
 def test_available_accelerators_in_registry():
-    assert AcceleratorRegistry.available_accelerators() == ["cpu", "cuda", "gpu", "hpu", "ipu", "mps", "tpu"]
+    assert AcceleratorRegistry.available_accelerators() == ["cpu", "cuda", "hpu", "ipu", "mps", "tpu"]

tests/tests_pytorch/lite/test_lite.py

@@ -315,9 +315,11 @@ def test_setup_dataloaders_replace_standard_sampler(shuffle, strategy):
     "accelerator, expected",
     [
         ("cpu", "cpu"),
+        pytest.param("cuda", "cuda:0", marks=RunIf(min_cuda_gpus=1)),
         pytest.param("gpu", "cuda:0", marks=RunIf(min_cuda_gpus=1)),
         pytest.param("tpu", "xla:0", marks=RunIf(tpu=True)),
         pytest.param("mps", "mps:0", marks=RunIf(mps=True)),
+        pytest.param("gpu", "mps:0", marks=RunIf(mps=True)),
     ],
 )
 def test_to_device(accelerator, expected):