Merge branch 'master' into loadams/auto-task-open-failure

README.md

@@ -225,6 +225,7 @@ Conduct](https://opensource.microsoft.com/codeofconduct/). For more information
 19. Sheng Shen, Zhewei Yao, Chunyuan Li, Trevor Darrell, Kurt Keutzer, Yuxiong He. (2023) Scaling Vision-Language Models with Sparse Mixture of Experts. [arXiv:2303.07226](https://arxiv.org/abs/2303.07226).
 20. Quentin Anthony, Ammar Ahmad Awan, Jeff Rasley, Yuxiong He, Aamir Shafi, Mustafa Abduljabbar, Hari Subramoni, Dhabaleswar Panda. (2023) MCR-DL: Mix-and-Match Communication Runtime for Deep Learning [arXiv:2303.08374](https://arxiv.org/abs/2303.08374) and will appear at IPDPS 2023.
 21. Siddharth Singh, Olatunji Ruwase, Ammar Ahmad Awan, Samyam Rajbhandari, Yuxiong He, Abhinav Bhatele. (2023) A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training [arXiv:2303.06318](https://arxiv.org/abs/2303.06318) and will appear at ICS 2023.
+22. Guanhua Wang, Heyang Qin, Sam Ade Jacobs, Connor Holmes, Samyam Rajbhandari, Olatunji Ruwase, Feng Yan, Lei Yang, Yuxiong He. (2023) ZeRO++: Extremely Efficient Collective Communication for Giant Model Training [arXiv:2306.10209](https://arxiv.org/abs/2306.10209).
 
 
 # Videos

accelerator/cpu_accelerator.py

@@ -40,7 +40,7 @@ class CPU_Accelerator(DeepSpeedAccelerator):
     def device_count(self):
         device_count = int(os.environ.get('LOCAL_SIZE', 0))
         if device_count > 0:
-            return os.environ.get('LOCAL_SIZE')
+            return device_count
         else:
             from deepspeed.utils.numa import get_numa_cores
             # Count NUMA node for number of cpu accelerators. On machine with HBM

accelerator/mps_accelerator.py

+# 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

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 @@ my_accelerator.communication_backend='nccl'
 my_accelerator.HalfTensor().device=device(type='cuda', index=0)
 my_accelerator.total_memory()=34089730048
 ---[output] python test_set.py---------
-'''
+"""

docs/index.md

@@ -126,6 +126,8 @@ comments.
 18. Syed Zawad, Cheng Li, Zhewei Yao, Elton Zheng, Yuxiong He, Feng Yan. (2023) DySR: Adaptive Super-Resolution via Algorithm and System Co-design. [ICLR:2023](https://openreview.net/forum?id=Pgtn4l6eKjv).
 19. Sheng Shen, Zhewei Yao, Chunyuan Li, Trevor Darrell, Kurt Keutzer, Yuxiong He. (2023) Scaling Vision-Language Models with Sparse Mixture of Experts. [arXiv:2303.07226](https://arxiv.org/abs/2303.07226).
 20. Quentin Anthony, Ammar Ahmad Awan, Jeff Rasley, Yuxiong He, Aamir Shafi, Mustafa Abduljabbar, Hari Subramoni, Dhabaleswar Panda. (2023) MCR-DL: Mix-and-Match Communication Runtime for Deep Learning [arXiv:2303.08374](https://arxiv.org/abs/2303.08374) and will appear at IPDPS 2023.
+21. Siddharth Singh, Olatunji Ruwase, Ammar Ahmad Awan, Samyam Rajbhandari, Yuxiong He, Abhinav Bhatele. (2023) A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training [arXiv:2303.06318](https://arxiv.org/abs/2303.06318) and will appear at ICS 2023.
+22. Guanhua Wang, Heyang Qin, Sam Ade Jacobs, Connor Holmes, Samyam Rajbhandari, Olatunji Ruwase, Feng Yan, Lei Yang, Yuxiong He. (2023) ZeRO++: Extremely Efficient Collective Communication for Giant Model Training [arXiv:2306.10209](https://arxiv.org/abs/2306.10209).
 
 # Videos
 1. DeepSpeed KDD 2020 Tutorial

version.txt

-0.10.0
+0.10.1