Repository: microsoft/DeepSpeed

Issue:
[REQUEST] M1 Max support
**Is your feature request related to a problem? Please describe.**
I'm looking to do some fine tuning of GPT-J on a Macbook Pro M1 Max w/ 64GB RAM

**Describe the solution you'd like**
Given the CPU, GPU, and memory are significant for this chip it seems a reasonable target and would skip the hassle and cost management/budget asks for cloud based training.


Discussion:
We’ve definitely been watching if/when pytorch will support M1, it sounds like it’s planned though.

https://github.com/pytorch/pytorch/issues/47702

Specifically see this comment: 
https://github.com/pytorch/pytorch/issues/47702#issuecomment-965625139

In terms of DeepSpeed support for M1 I suspect (depending on final design on torch side) many of our features will work well. However, we’ll have to reassess once torch releases their updated plan and final support here.
Looks like pytorch released support.
any updates?
+1 on the request
+1 for this, would be very helpful to have support from deepspeed in addition to torch.
+1 on the update, it would be very cool to have deepspeed now working on Apple M1/M2 machines now that PyTorch is supporting them..
It is exciting to see PyTorch support for M1/M2. We very open to extending deepspeed support to more and more accelerators but we currently lack bandwidth and hardware to explore this at the moment. However, we would gladly support any PRs for this similar to our ongoing support for the following Intel accelerator PR: https://github.com/microsoft/DeepSpeed/pull/2221
Any update? it's been a long time.
@phnessu4, unfortunately no update here as we have not had bandwidth or hardware access to drive this line. We would gladly accept any PR in this direction. Our accelerator abstraction is completed, and Intel XPU is now fully supported. 

Resolving patch:
diff --git a/accelerator/mps_accelerator.py b/accelerator/mps_accelerator.py
new file mode 100644
--- /dev/null
+++ b/accelerator/mps_accelerator.py
@@ -0,0 +1,225 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from .abstract_accelerator import DeepSpeedAccelerator
+
+# During setup stage torch may not be installed, pass on no torch will
+# allow op builder related API to be executed.
+try:
+    import torch.mps
+except ImportError:
+    pass
+
+
+class MPS_Accelerator(DeepSpeedAccelerator):
+
+    def __init__(self):
+        self._name = "mps"
+        self._communication_backend_name = None
+
+    def is_synchronized_device(self):
+        return False
+
+    # Device APIs
+    def device_name(self, device_index=None):
+        if device_index == None:
+            return "mps"
+        return "mps:{}".format(device_index)
+
+    def device(self, device_index):
+        return torch.device("mps", index=0)
+
+    def set_device(self, device_index):
+        return
+
+    def current_device(self):
+        return torch.device("mps", index=0)
+
+    def current_device_name(self):
+        return "mps:0"
+
+    def device_count(self):
+        return 1
+
+    def synchronize(self, device_index=None):
+        return torch.mps.synchronize()
+
+    # RNG APIs
+    def random(self):
+        return torch.random
+
+    def set_rng_state(self, new_state, device_index=None):
+        return torch.mps.set_rng_state(new_state)
+
+    def get_rng_state(self, device_index=None):
+        return torch.mps.get_rng_state()
+
+    def manual_seed(self, seed):
+        return torch.mps.manual_seed(seed)
+
+    def manual_seed_all(self, seed):
+        return torch.mps.manual_seed(seed)
+
+    def seed(self):
+        return torch.mps.seed()
+
+    def initial_seed(self, seed):
+        return
+
+    def default_generator(self, device_index):
+        return
+
+    # Streams/Events
+    @property
+    def Stream(self):
+        return None
+
+    def stream(self, stream):
+        return None
+
+    def current_stream(self, device_index=None):
+        return None
+
+    def default_stream(self, device_index=None):
+        return None
+
+    @property
+    def Event(self):
+        return None
+
+    # Memory management
+    def empty_cache(self):
+        return torch.mps.empty_cache()
+
+    def memory_allocated(self, device_index=None):
+        return torch.mps.current_allocated_memory()
+
+    def max_memory_allocated(self, device_index=None):
+        return torch.mps.driver_allocated_memory()
+
+    def set_per_process_memory_fraction(self, fraction):
+        return torch.mps.set_per_process_memory_fraction(fraction)
+
+    def reset_max_memory_allocated(self, device_index=None):
+        return
+
+    def memory_cached(self, device_index=None):
+        return
+
+    def max_memory_cached(self, device_index=None):
+        return
+
+    def reset_max_memory_cached(self, device_index=None):
+        return
+
+    def memory_stats(self, device_index=None):
+        return
+
+    def reset_peak_memory_stats(self, device_index=None):
+        return
+
+    def memory_reserved(self, device_index=None):
+        return
+
+    def max_memory_reserved(self, device_index=None):
+        return
+
+    def total_memory(self, device_index=None):
+        return
+
+    # Data types
+    def is_bf16_supported(self):
+        return False
+
+    def is_fp16_supported(self):
+        return False
+
+    # Misc
+    def amp(self):
+        return
+
+    def is_available(self):
+        return hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+
+    def range_push(self, msg):
+        return
+
+    def range_pop(self):
+        return
+
+    def lazy_call(self, callback):
+        return
+
+    def communication_backend_name(self):
+        return self._communication_backend_name
+
+    # Tensor operations
+    @property
+    def BFloat16Tensor(self):
+        return
+
+    @property
+    def ByteTensor(self):
+        return
+
+    @property
+    def DoubleTensor(self):
+        return
+
+    @property
+    def FloatTensor(self):
+        return
+
+    @property
+    def HalfTensor(self):
+        return
+
+    @property
+    def IntTensor(self):
+        return
+
+    @property
+    def LongTensor(self):
+        return
+
+    def pin_memory(self, tensor):
+        return tensor.pin_memory()
+
+    def on_accelerator(self, tensor):
+        device_str = str(tensor.device)
+        if device_str.startswith("mps"):
+            return True
+        else:
+            return False
+
+    def op_builder_dir(self):
+        try:
+            # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed
+            # if successful this also means we're doing a local install and not JIT compile path
+            from op_builder import __deepspeed__  # noqa: F401
+
+            return "op_builder"
+        except ImportError:
+            return "deepspeed.ops.op_builder"
+
+    # create an instance of op builder, specified by class_name
+    def create_op_builder(self, op_name):
+        builder_class = self.get_op_builder(op_name)
+        if builder_class != None:
+            return builder_class()
+        return None
+
+    # return an op builder class, specified by class_name
+    def get_op_builder(self, class_name):
+        from deepspeed.ops.op_builder.cpu import NotImplementedBuilder
+
+        return NotImplementedBuilder
+
+    def build_extension(self):
+        from torch.utils.cpp_extension import BuildExtension
+
+        return BuildExtension
diff --git a/accelerator/real_accelerator.py b/accelerator/real_accelerator.py
--- a/accelerator/real_accelerator.py
+++ b/accelerator/real_accelerator.py
@@ -35,10 +35,10 @@ def _validate_accelerator(accel_obj):
     # or deepspeed.accelerator.abstract_accelerator, consider accel_obj
     # is a conforming object
     if not ((dsa1 != None and isinstance(accel_obj, dsa1)) or (dsa2 != None and isinstance(accel_obj, dsa2))):
-        raise AssertionError(f'{accel_obj.__class__.__name__} accelerator is not subclass of DeepSpeedAccelerator')
+        raise AssertionError(f"{accel_obj.__class__.__name__} accelerator is not subclass of DeepSpeedAccelerator")
 
     # TODO: turn off is_available test since this breaks tests
-    #assert accel_obj.is_available(), \
+    # assert accel_obj.is_available(), \
     #    f'{accel_obj.__class__.__name__} accelerator fails is_available() test'
 
 
@@ -51,32 +51,41 @@ def get_accelerator():
     ds_set_method = None
     # 1. Detect whether there is override of DeepSpeed accelerators from environment variable.
     #    DS_ACCELERATOR = 'cuda'|'xpu'|'cpu'
-    if 'DS_ACCELERATOR' in os.environ.keys():
-        accelerator_name = os.environ['DS_ACCELERATOR']
-        if accelerator_name == 'xpu':
+    if "DS_ACCELERATOR" in os.environ.keys():
+        accelerator_name = os.environ["DS_ACCELERATOR"]
+        if accelerator_name == "xpu":
             try:
                 from intel_extension_for_deepspeed import XPU_Accelerator  # noqa: F401
             except ImportError as e:
                 raise ValueError(
-                    f'XPU_Accelerator requires intel_extension_for_deepspeed, which is not installed on this system.')
-        elif accelerator_name == 'cpu':
+                    f"XPU_Accelerator requires intel_extension_for_deepspeed, which is not installed on this system.")
+        elif accelerator_name == "cpu":
             try:
                 import intel_extension_for_pytorch  # noqa: F401
             except ImportError as e:
                 raise ValueError(
-                    f'CPU_Accelerator requires intel_extension_for_pytorch, which is not installed on this system.')
-        elif accelerator_name == 'cuda':
+                    f"CPU_Accelerator requires intel_extension_for_pytorch, which is not installed on this system.")
+        elif accelerator_name == "cuda":
             pass
+        elif accelerator_name == "mps":
+            try:
+                import torch.mps
+
+                # should use torch.mps.is_available() if it exists someday but this is used as proxy
+                torch.mps.current_allocated_memory()
+            except (RuntimeError, ImportError) as e:
+                raise ValueError(f"MPS_Accelerator requires torch.mps, which is not installed on this system.")
         else:
             raise ValueError(
                 f'DS_ACCELERATOR must be one of "cuda", "cpu", or "xpu".  Value "{accelerator_name}" is not supported')
-        ds_set_method = 'override'
+        ds_set_method = "override"
 
     # 2. If no override, detect which accelerator to use automatically
     if accelerator_name == None:
         try:
             from intel_extension_for_deepspeed import XPU_Accelerator  # noqa: F401,F811
-            accelerator_name = 'xpu'
+
+            accelerator_name = "xpu"
         except ImportError as e:
             # We need a way to choose between CUDA_Accelerator and CPU_Accelerator
             # Currently we detect whether intel_extension_for_pytorch is installed
@@ -90,21 +99,35 @@ def get_accelerator():
             #    between installation time and runtime.
             try:
                 import intel_extension_for_pytorch  # noqa: F401,F811
-                accelerator_name = 'cpu'
+
+                accelerator_name = "cpu"
             except ImportError as e:
-                accelerator_name = 'cuda'
-        ds_set_method = 'auto detect'
+                try:
+                    import torch.mps
+
+                    # should use torch.mps.is_available() if it exists someday but this is used as proxy
+                    torch.mps.current_allocated_memory()
+                    accelerator_name = "mps"
+                except (RuntimeError, ImportError) as e:
+                    accelerator_name = "cuda"
+        ds_set_method = "auto detect"
 
     # 3. Set ds_accelerator accordingly
-    if accelerator_name == 'cuda':
+    if accelerator_name == "cuda":
         from .cuda_accelerator import CUDA_Accelerator
+
         ds_accelerator = CUDA_Accelerator()
-    elif accelerator_name == 'cpu':
+    elif accelerator_name == "cpu":
         from .cpu_accelerator import CPU_Accelerator
+
         ds_accelerator = CPU_Accelerator()
-    elif accelerator_name == 'xpu':
+    elif accelerator_name == "xpu":
         # XPU_Accelerator is already imported in detection stage
         ds_accelerator = XPU_Accelerator()
+    elif accelerator_name == "mps":
+        from .mps_accelerator import MPS_Accelerator
+
+        ds_accelerator = MPS_Accelerator()
     _validate_accelerator(ds_accelerator)
     if accel_logger is not None:
         accel_logger.info(f"Setting ds_accelerator to {ds_accelerator._name} ({ds_set_method})")
@@ -119,7 +142,7 @@ def set_accelerator(accel_obj):
     ds_accelerator = accel_obj
 
 
-'''
+"""
 -----------[code] test_get.py -----------
 from deepspeed.accelerator import get_accelerator
 my_accelerator = get_accelerator()
@@ -161,4 +184,4 @@ def set_accelerator(accel_obj):
 my_accelerator.HalfTensor().device=device(type='cuda', index=0)
 my_accelerator.total_memory()=34089730048
 ---[output] python test_set.py---------
-'''
+"""


Test patch: