Merge pull request #13123 from Lightning-AI/mps_accelerator

MPS Accelerator

Author: justusschock

CHANGELOG.md

@@ -78,9 +78,11 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `XLAEnvironment` cluster environment plugin ([#11330](https://github.com/PyTorchLightning/pytorch-lightning/pull/11330))
 
 
+- Added support for calling unknown methods with `DummyLogger` ([#13224](https://github.com/PyTorchLightning/pytorch-lightning/pull/13224)
 
 
-- Added support for calling unknown methods with `DummyLogger` ([#13224](https://github.com/PyTorchLightning/pytorch-lightning/pull/13224))
+- Added Apple Silicon Support via `MPSAccelerator` ([#13123](https://github.com/PyTorchLightning/pytorch-lightning/pull/13123))
+
 
 ### Changed
 

docs/source-pytorch/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-pytorch/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::
+
+   Both the MPS accelerator and the PyTorch backend are still experimental.
+   As such, not all operations are currently supported. However, with ongoing development from the PyTorch team, an increasingly large number of operations are becoming available.
+   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 CPU-cores, GPU-cores, a neural engine and shared memory between all of these features.
+
+----
+
+So it's a CPU?
+--------------
+Apple silicon includes CPU-cores among several other features. However, the full potential for the hardware acceleration of which the M-Socs are capable is unavailable when running on the ``CPUAccelerator``. This is 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-pytorch/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-pytorch/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
 

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

src/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]:

src/pytorch_lightning/accelerators/mps.py

@@ -0,0 +1,95 @@
+# 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
+
+_MPS_AVAILABLE = _TORCH_GREATER_EQUAL_1_12 and torch.backends.mps.is_available()
+
+
+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,
+    }

src/pytorch_lightning/lite/lite.py

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

src/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 (
@@ -178,7 +179,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
@@ -407,7 +408,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(
@@ -484,6 +485,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"
@@ -571,11 +574,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: