MPS Accelerator

CHANGELOG.md

@@ -65,6 +65,13 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added `teardown()` method to `Accelerator` ([#11935](https://github.com/PyTorchLightning/pytorch-lightning/pull/11935))
 
+
+- Added Apple Silicon Support via `MPSAccelerator` ([#13123](https://github.com/PyTorchLightning/pytorch-lightning/pull/13123))
+
+
+-
+
+
 ### Changed
 
 - Enable validation during overfitting ([#12527](https://github.com/PyTorchLightning/pytorch-lightning/pull/12527))

docs/source/accelerators/mps.rst

@@ -0,0 +1,32 @@
+.. _mps:
+
+Accelerator: Apple Silicon training
+===================================
+
+.. raw:: html
+
+    <div class="display-card-container">
+        <div class="row">
+
+.. Add callout items below this line
+
+.. displayitem::
+   :header: Prepare your code (Optional)
+   :description: Prepare your code to run on any hardware
+   :col_css: col-md-4
+   :button_link: accelerator_prepare.html
+   :height: 150
+   :tag: basic
+
+.. displayitem::
+   :header: Basic
+   :description: Learn the basics of Apple silicon gpu training.
+   :col_css: col-md-4
+   :button_link: mps_basic.html
+   :height: 150
+   :tag: basic
+
+.. raw:: html
+
+        </div>
+    </div>

docs/source/accelerators/mps_basic.rst

@@ -0,0 +1,48 @@
+:orphan:
+
+.. _mps_basic:
+
+MPS training (basic)
+====================
+**Audience:** Users looking to train on their Apple silicon GPUs.
+
+.. warning::
+
+   The MPS accelerator as well as the PyTorch backend are still very experimental.
+   So far not all operations are supported, but more ops are coming every day due to development from the PyTorch Team.
+   You can use ``PYTORCH_ENABLE_MPS_FALLBACK=1 python your_script.py`` to fall back to cpu for unsupported operations.
+
+
+----
+
+What is Apple silicon?
+----------------------
+Apple silicon chips are a unified system on a chip (SoC) developed by Apple based on the ARM design.
+Among other things, they feature a CPU-cores, GPU-cores, a neural engine and shared memory between all of those.
+
+----
+
+So it's a CPU?
+--------------
+Among other things yes, it includes CPU-cores. However, when running on the ``CPUAccelerator``, not the full potential of hardware acceleration the M-Socs are capable of, is used because they also feature a GPU and a neural engine.
+
+To use them, Lightning supports the ``MPSAccelerator``.
+
+----
+
+Run on Apple silicon gpus
+-------------------------
+Enable the following Trainer arguments to run on Apple silicon gpus (MPS devices).
+
+.. code::
+
+   trainer = Trainer(accelerator="mps", devices=1)
+
+.. note::
+   The ``MPSAccelerator`` only supports 1 device at a time. Currently there are no machines with multiple MPS-capable GPUs.
+
+----
+
+What does MPS stand for?
+------------------------
+MPS is short for `Metal Performance Shaders <https://developer.apple.com/metal/>`_  which is the technology used in the back for gpu communication and computing.

docs/source/extensions/accelerator.rst

@@ -4,14 +4,15 @@
 Accelerator
 ###########
 
-The Accelerator connects a Lightning Trainer to arbitrary hardware (CPUs, GPUs, TPUs, IPUs, ...).
+The Accelerator connects a Lightning Trainer to arbitrary hardware (CPUs, GPUs, TPUs, IPUs, MPS, ...).
 Currently there are accelerators for:
 
 - CPU
 - :doc:`GPU <../accelerators/gpu>`
 - :doc:`TPU <../accelerators/tpu>`
 - :doc:`IPU <../accelerators/ipu>`
 - :doc:`HPU <../accelerators/hpu>`
+- :doc:`MPS <../accelerators/mps>`
 
 The Accelerator is part of the Strategy which manages communication across multiple devices (distributed communication).
 Whenever the Trainer, the loops or any other component in Lightning needs to talk to hardware, it calls into the Strategy and the Strategy calls into the Accelerator.
@@ -127,4 +128,5 @@ Accelerator API
     GPUAccelerator
     HPUAccelerator
     IPUAccelerator
+    MPSAccelerator
     TPUAccelerator

docs/source/index.rst

@@ -210,6 +210,7 @@ Current Lightning Users
    Train on single or multiple HPUs <accelerators/hpu>
    Train on single or multiple IPUs <accelerators/ipu>
    Train on single or multiple TPUs <accelerators/tpu>
+   Train on MPS <accelerators/mps>
    Use a pretrained model <advanced/pretrained>
    model/own_your_loop
 

pytorch_lightning/accelerators/__init__.py

@@ -15,6 +15,7 @@
 from pytorch_lightning.accelerators.gpu import GPUAccelerator  # noqa: F401
 from pytorch_lightning.accelerators.hpu import HPUAccelerator  # noqa: F401
 from pytorch_lightning.accelerators.ipu import IPUAccelerator  # noqa: F401
+from pytorch_lightning.accelerators.mps import MPSAccelerator  # noqa: F401
 from pytorch_lightning.accelerators.registry import AcceleratorRegistry, call_register_accelerators  # noqa: F401
 from pytorch_lightning.accelerators.tpu import TPUAccelerator  # noqa: F401
 

pytorch_lightning/accelerators/gpu.py

@@ -74,7 +74,7 @@ def get_device_stats(self, device: _DEVICE) -> Dict[str, Any]:
     @staticmethod
     def parse_devices(devices: Union[int, str, List[int]]) -> Optional[List[int]]:
         """Accelerator device parsing logic."""
-        return device_parser.parse_gpu_ids(devices)
+        return device_parser.parse_gpu_ids(devices, include_cuda=True)
 
     @staticmethod
     def get_parallel_devices(devices: List[int]) -> List[torch.device]:

pytorch_lightning/accelerators/mps.py

@@ -0,0 +1,98 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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, List, Optional, Union
+
+import torch
+
+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 _PSUTIL_AVAILABLE, _TORCH_GREATER_EQUAL_1_12
+from pytorch_lightning.utilities.types import _DEVICE
+
+if _TORCH_GREATER_EQUAL_1_12:
+    _MPS_AVAILABLE = torch.backends.mps.is_available()
+else:
+    _MPS_AVAILABLE = False
+
+
+class MPSAccelerator(Accelerator):
+    """Accelerator for Metal Apple Silicon GPU devices."""
+
+    def setup_environment(self, root_device: torch.device) -> None:
+        """
+        Raises:
+            MisconfigurationException:
+                If the selected device is not MPS.
+        """
+        super().setup_environment(root_device)
+        if root_device.type != "mps":
+            raise MisconfigurationException(f"Device should be MPS, got {root_device} instead.")
+
+    def get_device_stats(self, device: _DEVICE) -> Dict[str, Any]:
+        """Get M1 (cpu + gpu) stats from ``psutil`` package."""
+        return get_device_stats()
+
+    @staticmethod
+    def parse_devices(devices: Union[int, str, List[int]]) -> Optional[List[int]]:
+        """Accelerator device parsing logic."""
+        parsed_devices = device_parser.parse_gpu_ids(devices, include_mps=True)
+        return parsed_devices
+
+    @staticmethod
+    def get_parallel_devices(devices: Union[int, str, List[int]]) -> List[torch.device]:
+        """Gets parallel devices for the Accelerator."""
+        parsed_devices = MPSAccelerator.parse_devices(devices)
+        assert parsed_devices is not None
+
+        return [torch.device("mps", i) for i in range(len(parsed_devices))]
+
+    @staticmethod
+    def auto_device_count() -> int:
+        """Get the devices when set to auto."""
+        return 1
+
+    @staticmethod
+    def is_available() -> bool:
+        """MPS is only available for certain torch builds starting at torch>=1.12."""
+        return _MPS_AVAILABLE
+
+    @classmethod
+    def register_accelerators(cls, accelerator_registry: Dict) -> None:
+        accelerator_registry.register(
+            "mps",
+            cls,
+            description=cls.__class__.__name__,
+        )
+
+
+# device metrics
+_VM_PERCENT = "M1_vm_percent"
+_PERCENT = "M1_percent"
+_SWAP_PERCENT = "M1_swap_percent"
+
+
+def get_device_stats() -> Dict[str, float]:
+    if not _PSUTIL_AVAILABLE:
+        raise ModuleNotFoundError(
+            "Fetching M1 device stats requires `psutil` to be installed."
+            " Install it by running `pip install -U psutil`."
+        )
+    import psutil
+
+    return {
+        _VM_PERCENT: psutil.virtual_memory().percent,
+        _PERCENT: psutil.cpu_percent(),
+        _SWAP_PERCENT: psutil.swap_memory().percent,
+    }

pytorch_lightning/lite/lite.py

@@ -465,6 +465,7 @@ def _supported_device_types() -> Sequence[_AcceleratorType]:
             _AcceleratorType.CPU,
             _AcceleratorType.GPU,
             _AcceleratorType.TPU,
+            _AcceleratorType.MPS,
         )
 
     @staticmethod

pytorch_lightning/trainer/connectors/accelerator_connector.py

@@ -25,6 +25,7 @@
 from pytorch_lightning.accelerators.gpu import GPUAccelerator
 from pytorch_lightning.accelerators.hpu import HPUAccelerator
 from pytorch_lightning.accelerators.ipu import IPUAccelerator
+from pytorch_lightning.accelerators.mps import MPSAccelerator
 from pytorch_lightning.accelerators.registry import AcceleratorRegistry
 from pytorch_lightning.accelerators.tpu import TPUAccelerator
 from pytorch_lightning.plugins import (
@@ -173,7 +174,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._parallel_devices: List[Union[int, torch.device, str]] = []
         self._layer_sync: Optional[LayerSync] = NativeSyncBatchNorm() if sync_batchnorm else None
         self.checkpoint_io: Optional[CheckpointIO] = None
         self._amp_type_flag: Optional[LightningEnum] = None
@@ -402,7 +403,7 @@ def _check_device_config_and_set_final_flags(
         if self._devices_flag == "auto" and self._accelerator_flag is None:
             raise MisconfigurationException(
                 f"You passed `devices={devices}` but haven't specified"
-                " `accelerator=('auto'|'tpu'|'gpu'|'ipu'|'cpu'|'hpu)` for the devices mapping."
+                " `accelerator=('auto'|'tpu'|'gpu'|'ipu'|'cpu'|'hpu'|'mps')` for the devices mapping."
             )
 
     def _map_deprecated_devices_specific_info_to_accelerator_and_device_flag(
@@ -479,6 +480,8 @@ def _choose_accelerator(self) -> str:
                 return "ipu"
             if _HPU_AVAILABLE:
                 return "hpu"
+            if MPSAccelerator.is_available():
+                return "mps"
             if torch.cuda.is_available() and torch.cuda.device_count() > 0:
                 return "gpu"
         return "cpu"
@@ -566,11 +569,13 @@ def _choose_strategy(self) -> Union[Strategy, str]:
         if self._num_nodes_flag > 1:
             return DDPStrategy.strategy_name
         if len(self._parallel_devices) <= 1:
-            device = (
-                device_parser.determine_root_gpu_device(self._parallel_devices)  # type: ignore
-                if self._accelerator_flag == "gpu"
-                else "cpu"
-            )
+            # TODO: Change this once gpu accelerator was renamed to cuda accelerator
+            if isinstance(self._accelerator_flag, (GPUAccelerator, MPSAccelerator)) or (
+                isinstance(self._accelerator_flag, str) and self._accelerator_flag in ("gpu", "mps")
+            ):
+                device = device_parser.determine_root_gpu_device(self._parallel_devices)
+            else:
+                device = "cpu"
             # TODO: lazy initialized device, then here could be self._strategy_flag = "single_device"
             return SingleDeviceStrategy(device=device)  # type: ignore
         if len(self._parallel_devices) > 1:

pytorch_lightning/trainer/trainer.py

@@ -35,7 +35,14 @@
 from torch.utils.data import DataLoader
 
 import pytorch_lightning as pl
-from pytorch_lightning.accelerators import Accelerator, GPUAccelerator, HPUAccelerator, IPUAccelerator, TPUAccelerator
+from pytorch_lightning.accelerators import (
+    Accelerator,
+    GPUAccelerator,
+    HPUAccelerator,
+    IPUAccelerator,
+    MPSAccelerator,
+    TPUAccelerator,
+)
 from pytorch_lightning.callbacks import Callback, EarlyStopping, ModelCheckpoint, ProgressBarBase
 from pytorch_lightning.callbacks.prediction_writer import BasePredictionWriter
 from pytorch_lightning.core.datamodule import LightningDataModule
@@ -189,7 +196,7 @@ def __init__(
 
         Args:
 
-            accelerator: Supports passing different accelerator types ("cpu", "gpu", "tpu", "ipu", "hpu", "auto")
+            accelerator: Supports passing different accelerator types ("cpu", "gpu", "tpu", "ipu", "hpu", "mps, "auto")
                 as well as custom accelerator instances.
 
                 .. deprecated:: v1.5
@@ -1745,9 +1752,19 @@ def __setup_profiler(self) -> None:
         self.profiler.setup(stage=self.state.fn._setup_fn, local_rank=local_rank, log_dir=self.log_dir)
 
     def _log_device_info(self) -> None:
-        rank_zero_info(
-            f"GPU available: {torch.cuda.is_available()}, used: {isinstance(self.accelerator, GPUAccelerator)}"
-        )
+
+        if GPUAccelerator.is_available():
+            gpu_available = True
+            gpu_type = " (cuda)"
+        elif MPSAccelerator.is_available():
+            gpu_available = True
+            gpu_type = " (mps)"
+        else:
+            gpu_available = False
+            gpu_type = ""
+
+        gpu_used = isinstance(self.accelerator, (GPUAccelerator, MPSAccelerator))
+        rank_zero_info(f"GPU available: {gpu_available}{gpu_type}, used: {gpu_used}")
 
         num_tpu_cores = self.num_devices if isinstance(self.accelerator, TPUAccelerator) else 0
         rank_zero_info(f"TPU available: {_TPU_AVAILABLE}, using: {num_tpu_cores} TPU cores")
@@ -1758,6 +1775,7 @@ def _log_device_info(self) -> None:
         num_hpus = self.num_devices if isinstance(self.accelerator, HPUAccelerator) else 0
         rank_zero_info(f"HPU available: {_HPU_AVAILABLE}, using: {num_hpus} HPUs")
 
+        # TODO: Integrate MPS Accelerator here, once gpu maps to both
         if torch.cuda.is_available() and not isinstance(self.accelerator, GPUAccelerator):
             rank_zero_warn(
                 "GPU available but not used. Set `accelerator` and `devices` using"
@@ -1783,6 +1801,12 @@ def _log_device_info(self) -> None:
                 f" `Trainer(accelerator='hpu', devices={HPUAccelerator.auto_device_count()})`."
             )
 
+        if MPSAccelerator.is_available() and not isinstance(self.accelerator, MPSAccelerator):
+            rank_zero_warn(
+                "MPS available but not used. Set `accelerator` and `devices` using"
+                f" `Trainer(accelerator='mps', devices={MPSAccelerator.auto_device_count()})`."
+            )
+
     """
     Data loading methods
     """