Merge branch 'master' into tohtana/fix-zero2-bf16-shared-grad-8224

Author: sfc-gh-truwase

accelerator/abstract_accelerator.py

@@ -174,6 +174,10 @@ def is_bf16_supported(self):
     def is_fp16_supported(self):
         ...
 
+    # Not abstract: nearly every accelerator supports fp64, so only those that do not need to override.
+    def is_fp64_supported(self):
+        return True
+
     @abc.abstractmethod
     def supported_dtypes(self):
         ...

accelerator/cuda_accelerator.py

@@ -274,7 +274,7 @@ def prefer_triton_grouped_mm(self):
         # torch._grouped_mm only has a fused grouped-GEMM kernel on Hopper (sm90)
         # and newer; on sm8x it falls back to a slow per-group loop, so a Triton
         # grouped-GEMM kernel is preferred there when Triton is available.
-        from deepspeed.moe.group_gemm_triton import is_available as triton_grouped_mm_is_available
+        from deepspeed.ops.triton_ops import is_triton_available as triton_grouped_mm_is_available
         # not verified on AMD GPU
         if torch.version.hip is not None or not triton_grouped_mm_is_available():
             return False

accelerator/mps_accelerator.py

@@ -19,11 +19,14 @@ class MPS_Accelerator(DeepSpeedAccelerator):
 
     def __init__(self):
         self._name = "mps"
-        self._communication_backend_name = None
+        # MPS has no native collective backend; gloo is the only torch backend available on macOS.
+        self._communication_backend_name = "gloo"
         self._compile_backend = "inductor"
 
     def is_synchronized_device(self):
-        return False
+        # MPS runs everything on a single in-order command queue and exposes no user-visible
+        # streams, so DeepSpeed never needs to synchronize between streams on this device.
+        return True
 
     def use_host_timers(self):
         # Event timers are not supported on MPS
@@ -90,7 +93,8 @@ def Stream(self):
         return None
 
     def stream(self, stream):
-        return None
+        from deepspeed.runtime.utils import noop_context
+        return noop_context()
 
     def current_stream(self, device_index=None):
         return None
@@ -100,7 +104,7 @@ def default_stream(self, device_index=None):
 
     @property
     def Event(self):
-        return None
+        return torch.mps.Event
 
     # Memory management
     def empty_cache(self):
@@ -119,41 +123,53 @@ def reset_max_memory_allocated(self, device_index=None):
         return
 
     def memory_cached(self, device_index=None):
-        return
+        return torch.mps.driver_allocated_memory()
 
     def max_memory_cached(self, device_index=None):
-        return
+        return torch.mps.driver_allocated_memory()
 
     def reset_max_memory_cached(self, device_index=None):
         return
 
     def memory_stats(self, device_index=None):
-        return
+        # torch.mps has no caching-allocator stats; expose what it does report under the CUDA key names.
+        return {
+            'allocated_bytes.all.current': torch.mps.current_allocated_memory(),
+            'reserved_bytes.all.current': torch.mps.driver_allocated_memory(),
+        }
 
     def reset_peak_memory_stats(self, device_index=None):
         return
 
     def memory_reserved(self, device_index=None):
-        return
+        return torch.mps.driver_allocated_memory()
 
     def max_memory_reserved(self, device_index=None):
-        return
+        return torch.mps.driver_allocated_memory()
 
     def total_memory(self, device_index=None):
-        return
+        # Unified memory: the driver-recommended working set is the usable budget for the GPU.
+        return torch.mps.recommended_max_memory()
 
     def available_memory(self, device_index=None):
-        return
+        return self.total_memory() - torch.mps.driver_allocated_memory()
 
     # Data types
     def is_bf16_supported(self):
-        return False
+        # bf16 on MPS requires macOS 14 (Sonoma) or newer.
+        return torch.backends.mps.is_macos_or_newer(14, 0)
 
     def is_fp16_supported(self):
+        return True
+
+    def is_fp64_supported(self):
         return False
 
     def supported_dtypes(self):
-        return [torch.float]
+        supported_dtypes = [torch.float, torch.half]
+        if self.is_bf16_supported():
+            supported_dtypes.append(torch.bfloat16)
+        return supported_dtypes
 
     # Misc
     def is_available(self):
@@ -188,31 +204,39 @@ def replay_graph(self, graph):
     # Tensor operations
     @property
     def BFloat16Tensor(self):
-        return
+        return torch.BFloat16Tensor
 
     @property
     def ByteTensor(self):
-        return
+        return torch.ByteTensor
 
     @property
     def DoubleTensor(self):
-        return
+        return torch.DoubleTensor
 
     @property
     def FloatTensor(self):
-        return
+        return torch.FloatTensor
 
     @property
     def HalfTensor(self):
-        return
+        return torch.HalfTensor
 
     @property
     def IntTensor(self):
-        return
+        return torch.IntTensor
 
     @property
     def LongTensor(self):
-        return
+        return torch.LongTensor
+
+    # Apple Silicon has unified memory, so host tensors are already directly accessible to the GPU
+    # and there is nothing to pin. torch's pin_memory() also raises for the MPS backend.
+    def _torch_pin_memory(self, tensor):
+        return tensor
+
+    def _torch_is_pinned(self, tensor):
+        return tensor.device.type == 'cpu'
 
     def on_accelerator(self, tensor):
         device_str = str(tensor.device)
@@ -227,9 +251,9 @@ def op_builder_dir(self):
             # if successful this also means we're doing a local install and not JIT compile path
             from op_builder import __deepspeed__  # noqa: F401 # type: ignore
 
-            return "op_builder"
+            return "op_builder.mps"
         except ImportError:
-            return "deepspeed.ops.op_builder"
+            return "deepspeed.ops.op_builder.mps"
 
     # create an instance of op builder, specified by class_name
     def create_op_builder(self, op_name):
@@ -240,9 +264,18 @@ def create_op_builder(self, op_name):
 
     # return an op builder class, specified by class_name
     def get_op_builder(self, class_name):
-        from deepspeed.ops.op_builder.cpu import NotImplementedBuilder
+        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 # type: ignore
+            from op_builder.mps import FusedAdamBuilder, NotImplementedBuilder
+        except ImportError:
+            from deepspeed.ops.op_builder.mps import FusedAdamBuilder, NotImplementedBuilder
 
-        return NotImplementedBuilder
+        if class_name == "FusedAdamBuilder":
+            return FusedAdamBuilder
+        else:
+            return NotImplementedBuilder
 
     def build_extension(self):
         from torch.utils.cpp_extension import BuildExtension

deepspeed/comm/torch.py

@@ -95,6 +95,66 @@ def wait(self):
         return None
 
 
+class StagedWork:
+    """Completes a staged collective by copying the CPU results back to the original device tensors."""
+
+    def __init__(self, work, copy_back):
+        self.work = work
+        self.copy_back = copy_back
+
+    def wait(self):
+        if self.work is not None:
+            self.work.wait()
+        self.copy_back()
+        return None
+
+
+def _needs_cpu_staging(tensor):
+    # gloo (the only torch backend on macOS) cannot operate on MPS tensors.
+    return isinstance(tensor, torch.Tensor) and tensor.device.type == 'mps'
+
+
+def stage_on_cpu(func):
+    """Runs a collective on CPU copies of any MPS tensor arguments, then copies the results back.
+
+    This is what lets DeepSpeed use the gloo backend on Apple Silicon, where device tensors are
+    not supported by any torch.distributed backend. Unified memory keeps the copies cheap.
+    """
+
+    def _stage(arg, pairs):
+        if _needs_cpu_staging(arg):
+            cpu_tensor = arg.to('cpu')
+            pairs.append((arg, cpu_tensor))
+            return cpu_tensor
+        if isinstance(arg, list) and any(_needs_cpu_staging(t) for t in arg):
+            return [_stage(t, pairs) for t in arg]
+        return arg
+
+    signature = inspect.signature(func)
+
+    def wrapper(self, *args, **kwargs):
+        pairs = []
+        args = [_stage(arg, pairs) for arg in args]
+        kwargs = {key: _stage(arg, pairs) for key, arg in kwargs.items()}
+        if not pairs:
+            return func(self, *args, **kwargs)
+
+        work = func(self, *args, **kwargs)
+
+        def copy_back():
+            for device_tensor, cpu_tensor in pairs:
+                device_tensor.copy_(cpu_tensor)
+
+        # async_op is usually forwarded positionally, so resolve it against the real signature.
+        bound_args = signature.bind(self, *args, **kwargs)
+        if bound_args.arguments.get('async_op', False):
+            return StagedWork(work, copy_back)
+        copy_back()
+        return work
+
+    return wrapper
+
+
 class TorchBackend(Backend):
     """
         A light-weight wrapper class for torch.distributed API.
@@ -179,6 +239,7 @@ def init_process_group(self, backend, timeout, init_method, rank, world_size):
         self.using_mpi = torch.distributed.get_backend() == 'mpi'
 
     @disable_compiler_collective
+    @stage_on_cpu
     def all_reduce(self, tensor, op=torch.distributed.ReduceOp.SUM, group=None, async_op=False):
         op = self._reduce_op(op)
         return torch.distributed.all_reduce(tensor=tensor, op=op, group=group, async_op=async_op)
@@ -195,6 +256,7 @@ def inference_all_reduce(self, tensor, op, group=None):
             return torch.ops.deepspeed.inference_all_reduce_(tensor)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def all_reduce_coalesced(self, tensors, op=torch.distributed.ReduceOp.SUM, group=None, async_op=False):
         """ proxy func to torch.distributed.all_reduce_coalesced,
         which is included in PyTorch 1.13 and above
@@ -206,6 +268,7 @@ def all_reduce_coalesced(self, tensors, op=torch.distributed.ReduceOp.SUM, group
         return torch.distributed.all_reduce_coalesced(tensors=tensors, op=op, group=group, async_op=async_op)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def reduce(self, tensor, dst, op=ReduceOp.SUM, group=None, async_op=False):
         if DS_COMM_REDUCE_OFF:
             if int(os.getenv('RANK', '0')) == 0:
@@ -214,6 +277,7 @@ def reduce(self, tensor, dst, op=ReduceOp.SUM, group=None, async_op=False):
         return torch.distributed.reduce(tensor=tensor, dst=dst, op=self._reduce_op(op), group=group, async_op=async_op)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def reduce_scatter(self, output, input_list, op=ReduceOp.SUM, group=None, async_op=False):
         if DS_COMM_REDUCE_SCATTER_OFF:
             if int(os.getenv('RANK', '0')) == 0:
@@ -227,6 +291,7 @@ def reduce_scatter(self, output, input_list, op=ReduceOp.SUM, group=None, async_
                                                     async_op=async_op)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def broadcast(self, tensor, src, group=None, async_op=False):
         if DS_COMM_BROADCAST_OFF:
             if int(os.getenv('RANK', '0')) == 0:
@@ -240,6 +305,7 @@ def broadcast_object_list(self, object_list, src, group=None, device=None):
         return torch.distributed.broadcast_object_list(object_list=object_list, src=src, group=group, device=device)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def all_gather(self, tensor_list, tensor, group=None, async_op=False):
         if DS_COMM_ALL_GATHER_OFF:
             if int(os.getenv('RANK', '0')) == 0:
@@ -249,6 +315,7 @@ def all_gather(self, tensor_list, tensor, group=None, async_op=False):
             return torch.distributed.all_gather(tensor_list=tensor_list, tensor=tensor, group=group, async_op=async_op)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def all_gather_into_tensor(self, output_tensor, input_tensor, group=None, async_op=False):
         # Transparent SDMA fast-path on AMD/ROCm: when the mori backend is
         # available and the call is on the WORLD process group, route
@@ -266,6 +333,7 @@ def all_gather_into_tensor(self, output_tensor, input_tensor, group=None, async_
                                             async_op=async_op)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def all_gather_base(self, output_tensor, input_tensor, group=None, async_op=False):
         if DS_COMM_ALL_GATHER_OFF:
             if int(os.getenv('RANK', '0')) == 0:
@@ -284,6 +352,7 @@ def all_gather_base(self, output_tensor, input_tensor, group=None, async_op=Fals
                 pass
 
     @disable_compiler_collective
+    @stage_on_cpu
     def all_gather_coalesced(self, output_tensors, input_tensors, group=None, async_op=False):
         """"""
         assert len(output_tensors) == len(input_tensors), ""
@@ -312,6 +381,7 @@ def all_gather_object(self, object_list, obj, group=None):
         return torch.distributed.all_gather_object(object_list=object_list, obj=obj, group=group)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def reduce_scatter_tensor(self, output_tensor, input_tensor, op=ReduceOp.SUM, group=None, async_op=False):
         if self.has_reduce_scatter_tensor():
             return self.reduce_scatter_function(output_tensor,
@@ -326,6 +396,7 @@ def reduce_scatter_tensor(self, output_tensor, input_tensor, op=ReduceOp.SUM, gr
             pass
 
     @disable_compiler_collective
+    @stage_on_cpu
     def all_to_all_single(self,
                           output,
                           input,
@@ -341,26 +412,32 @@ def all_to_all_single(self,
                                                    async_op=async_op)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def all_to_all(self, output_tensor_list, input_tensor_list, group=None, async_op=False):
         return torch.distributed.all_to_all(output_tensor_list, input_tensor_list, group=group, async_op=async_op)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def send(self, tensor, dst, group=None, tag=0):
         return torch.distributed.send(tensor=tensor, dst=dst, group=group, tag=tag)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def recv(self, tensor, src=None, group=None, tag=0):
         return torch.distributed.recv(tensor=tensor, src=src, group=group, tag=tag)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def isend(self, tensor, dst, group=None, tag=0):
         return torch.distributed.isend(tensor=tensor, dst=dst, group=group, tag=tag)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def irecv(self, tensor, src=None, group=None, tag=0):
         return torch.distributed.irecv(tensor=tensor, src=src, group=group, tag=tag)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def gather(self, tensor, gather_list=None, dst=0, group=None, async_op=False):
         return torch.distributed.gather(tensor=tensor,
                                         gather_list=gather_list,
@@ -369,6 +446,7 @@ def gather(self, tensor, gather_list=None, dst=0, group=None, async_op=False):
                                         async_op=async_op)
 
     @disable_compiler_collective
+    @stage_on_cpu
     def scatter(self, tensor, scatter_list=None, src=0, group=None, async_op=False):
         return torch.distributed.scatter(tensor=tensor,
                                          scatter_list=scatter_list,