[Apple Silicon Support Phase 1] Add Metal FusedAdam kernel and CPU Adam build for Apple Silicon

MANIFEST.in

@@ -3,7 +3,7 @@ include deepspeed/inference/v2/kernels/ragged_ops/libs/*.so
 include deepspeed/inference/v2/kernels/cutlass_ops/libs/*.so
 recursive-include requirements *.txt
 recursive-include deepspeed *.cpp *.h *.hpp *.cu *.hip *.tr *.cuh *.cc *.json
-recursive-include csrc *.cpp *.h *.hpp *.cu *.tr *.cuh *.cc
+recursive-include csrc *.cpp *.h *.hpp *.cu *.tr *.cuh *.cc *.metal
 recursive-include op_builder *.py
 recursive-include benchmarks *.py
 recursive-include accelerator *.py

accelerator/mps_accelerator.py

@@ -268,12 +268,14 @@ def get_op_builder(self, class_name):
             # 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
+            from op_builder.mps import CPUAdamBuilder, FusedAdamBuilder, NotImplementedBuilder
         except ImportError:
-            from deepspeed.ops.op_builder.mps import FusedAdamBuilder, NotImplementedBuilder
+            from deepspeed.ops.op_builder.mps import CPUAdamBuilder, FusedAdamBuilder, NotImplementedBuilder
 
         if class_name == "FusedAdamBuilder":
             return FusedAdamBuilder
+        elif class_name == "CPUAdamBuilder":
+            return CPUAdamBuilder
         else:
             return NotImplementedBuilder
 

csrc/mps/fused_adam.metal

@@ -0,0 +1,58 @@
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+// Fused Adam / AdamW step for one tensor. Mirrors csrc/adam/multi_tensor_adam.cu: all math is done in
+// fp32 regardless of the storage dtype, so fp16/bf16 parameters see the same numerics as on CUDA.
+
+#include <metal_stdlib>
+using namespace metal;
+
+template <typename T>
+kernel void fused_adam(device T* param [[buffer(0)]],
+                       device const T* grad [[buffer(1)]],
+                       device T* exp_avg [[buffer(2)]],
+                       device T* exp_avg_sq [[buffer(3)]],
+                       constant float& lr [[buffer(4)]],
+                       constant float& beta1 [[buffer(5)]],
+                       constant float& beta2 [[buffer(6)]],
+                       constant float& epsilon [[buffer(7)]],
+                       constant float& weight_decay [[buffer(8)]],
+                       constant float& bias_correction1 [[buffer(9)]],
+                       constant float& bias_correction2 [[buffer(10)]],
+                       constant float& adam_w_mode [[buffer(11)]],
+                       uint i [[thread_position_in_grid]])
+{
+    float g = float(grad[i]);
+    float p = float(param[i]);
+    float m = float(exp_avg[i]);
+    float v = float(exp_avg_sq[i]);
+
+    // L2 mode folds weight decay into the gradient; AdamW mode applies it to the parameter.
+    if (adam_w_mode == 0.0f) { g += weight_decay * p; }
+    m = beta1 * m + (1.0f - beta1) * g;
+    v = beta2 * v + (1.0f - beta2) * g * g;
+    float denom = sqrt(v / bias_correction2) + epsilon;
+    float update = (m / bias_correction1) / denom;
+    if (adam_w_mode != 0.0f) { update += weight_decay * p; }
+
+    param[i] = T(p - lr * update);
+    exp_avg[i] = T(m);
+    exp_avg_sq[i] = T(v);
+}
+
+template [[host_name("fused_adam_float")]] kernel void fused_adam<float>(
+    device float*, device const float*, device float*, device float*,
+    constant float&, constant float&, constant float&, constant float&, constant float&,
+    constant float&, constant float&, constant float&, uint);
+template [[host_name("fused_adam_half")]] kernel void fused_adam<half>(
+    device half*, device const half*, device half*, device half*,
+    constant float&, constant float&, constant float&, constant float&, constant float&,
+    constant float&, constant float&, constant float&, uint);
+// bfloat is a Metal 3.1 type (macOS 14); older toolchains still get the float and half kernels.
+#if __METAL_VERSION__ >= 310
+template [[host_name("fused_adam_bfloat")]] kernel void fused_adam<bfloat>(
+    device bfloat*, device const bfloat*, device bfloat*, device bfloat*,
+    constant float&, constant float&, constant float&, constant float&, constant float&,
+    constant float&, constant float&, constant float&, uint);
+#endif

docs/_tutorials/accelerator-setup-guide.md

@@ -306,7 +306,9 @@ Launch a single-process job as usual; no hostfile is needed:
 ```
 deepspeed --num_gpus 1 train.py --deepspeed --deepspeed_config ds_config.json
 ```
-ZeRO stages 0 through 3 are supported with fp32, fp16, and bf16 (bf16 requires macOS 14 or newer). The fused Adam optimizer runs as a PyTorch implementation on MPS; ZeRO-Offload (`DeepSpeedCPUAdam`) is not yet available on this backend.
+ZeRO stages 0 through 3 are supported with fp32, fp16, and bf16 (bf16 requires macOS 14 or newer), with or without ZeRO-Offload.
+
+The fused Adam optimizer is a Metal kernel compiled at first use through `torch.mps.compile_shader`; no Xcode project or C++ build is involved. ZeRO-Offload uses the C++ `DeepSpeedCPUAdam` kernel, which is built just-in-time with the system clang. Apple's clang has no OpenMP, so the kernel is single-threaded unless Homebrew's `libomp` is installed (`brew install libomp`), in which case it is picked up automatically. Because Apple Silicon has unified memory, offloading to the CPU optimizer does not copy parameters between separate memories.
 
 ## Limitations
 * PyTorch exposes one MPS device per machine, so `device_count()` is 1 and multi-device data parallelism on a single Mac is not possible.

op_builder/mps/__init__.py

@@ -2,5 +2,6 @@
 
 # DeepSpeed Team
 
+from .cpu_adam import CPUAdamBuilder
 from .fused_adam import FusedAdamBuilder
 from .no_impl import NotImplementedBuilder

op_builder/mps/builder.py

@@ -2,6 +2,8 @@
 
 # DeepSpeed Team
 
+import os
+
 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
@@ -12,11 +14,7 @@
 
 
 class MPSOpBuilder(OpBuilder):
-    """Base class for ops on Apple Silicon.
-
-    Ops here are currently pure PyTorch (torch.mps) implementations, so there is nothing to compile.
-    Metal kernels will plug into the same builder classes later.
-    """
+    """Base class for ops on Apple Silicon that need no C++ compilation (pure torch.mps implementations)."""
 
     def sources(self):
         return []
@@ -27,3 +25,30 @@ def include_paths(self):
     def is_compatible(self, verbose=False):
         import torch
         return hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+
+
+class MetalOpBuilder(MPSOpBuilder):
+    """Base class for ops implemented as Metal shaders.
+
+    Shaders are compiled at load time through torch.mps.compile_shader, which dispatches them on
+    PyTorch's own MPS command stream, so no Xcode toolchain or C++ extension build is needed.
+    Subclasses list their .metal files in metal_sources() and get back the compiled library.
+    """
+    _libraries = {}
+
+    def metal_sources(self):
+        return []
+
+    def is_compatible(self, verbose=False):
+        import torch
+        return super().is_compatible(verbose) and hasattr(torch.mps, "compile_shader")
+
+    def load_metal_library(self):
+        import torch
+        if self.name not in MetalOpBuilder._libraries:
+            source = ""
+            for path in self.metal_sources():
+                with open(os.path.join(self.deepspeed_src_path(path))) as shader_file:
+                    source += shader_file.read() + "\n"
+            MetalOpBuilder._libraries[self.name] = torch.mps.compile_shader(source)
+        return MetalOpBuilder._libraries[self.name]

op_builder/mps/cpu_adam.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import subprocess
+
+from .builder import MPSOpBuilder
+
+
+class CPUAdamBuilder(MPSOpBuilder):
+    """Builds the C++ CPU Adam kernel on Apple Silicon for ZeRO-Offload.
+
+    Unified memory means the optimizer step can run on the CPU cores without copying parameters
+    off the GPU, so this is the offload path on Macs until a Metal Adam kernel exists.
+    """
+    BUILD_VAR = "DS_BUILD_CPU_ADAM"
+    NAME = "cpu_adam"
+
+    def __init__(self):
+        super().__init__(name=self.NAME)
+
+    def absolute_name(self):
+        return f'deepspeed.ops.adam.{self.NAME}_op'
+
+    def sources(self):
+        return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp']
+
+    def include_paths(self):
+        return ['csrc/includes']
+
+    def builder(self):
+        from torch.utils.cpp_extension import CppExtension
+        include_dirs = [os.path.abspath(path) for path in self.strip_empty_entries(self.include_paths())]
+        return CppExtension(name=self.absolute_name(),
+                            sources=self.strip_empty_entries(self.sources()),
+                            include_dirs=include_dirs,
+                            extra_compile_args={'cxx': self.strip_empty_entries(self.cxx_args())},
+                            extra_link_args=self.strip_empty_entries(self.extra_ldflags()))
+
+    def _libomp_prefix(self):
+        # Apple clang ships without OpenMP; Homebrew's libomp provides it when installed.
+        try:
+            return subprocess.check_output(['brew', '--prefix', 'libomp'], stderr=subprocess.DEVNULL).decode().strip()
+        except (OSError, subprocess.CalledProcessError):
+            return None
+
+    def cxx_args(self):
+        args = ['-O3', '-std=c++17', '-g', '-Wno-reorder', '-D__SCALAR__']
+        libomp = self._libomp_prefix()
+        if libomp is not None:
+            args += ['-Xpreprocessor', '-fopenmp', f'-I{libomp}/include']
+        return args
+
+    def extra_ldflags(self):
+        libomp = self._libomp_prefix()
+        if libomp is None:
+            return []
+        return [f'-L{libomp}/lib', '-lomp']

op_builder/mps/fused_adam.py

@@ -2,7 +2,7 @@
 
 # DeepSpeed Team
 
-from .builder import MPSOpBuilder
+from .builder import MPSOpBuilder, MetalOpBuilder
 
 try:
     import torch
@@ -11,28 +11,73 @@
 
 
 class MPSFusedAdam:
-    """Pure-torch replacement for the CUDA multi_tensor_adam kernel, using torch._foreach_* ops on MPS.
+    """Drop-in for the CUDA multi_tensor_adam op: one Metal launch per tensor, math in fp32.
 
-    The math mirrors csrc/adam/multi_tensor_adam.cu so checkpoints and numerics stay interchangeable.
+    Falls back to torch._foreach_* ops when torch.mps.compile_shader is unavailable.
     """
+    kernels = None
+    compile_failed = False
 
     @staticmethod
     def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode,
                           bias_correction, weight_decay, *args):
+        bias_correction1 = 1.0
+        bias_correction2 = 1.0
+        if bias_correction:
+            bias_correction1 = 1.0 - beta1**step
+            bias_correction2 = 1.0 - beta2**step
+
         # The caller passes bf16 params as leaf tensors (not .data), so the in-place update must bypass autograd.
         with torch.no_grad():
-            MPSFusedAdam._adam_step(tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, bias_correction,
-                                    weight_decay)
+            if MPSFusedAdam.kernels is None:
+                MPSFusedAdam._foreach_adam_step(tensor_lists, lr, beta1, beta2, epsilon, adam_w_mode, bias_correction1,
+                                                bias_correction2, weight_decay)
+            else:
+                MPSFusedAdam._metal_adam_step(tensor_lists, lr, beta1, beta2, epsilon, adam_w_mode, bias_correction1,
+                                              bias_correction2, weight_decay)
 
     @staticmethod
-    def _adam_step(tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, bias_correction, weight_decay):
+    def _metal_adam_step(tensor_lists, lr, beta1, beta2, epsilon, adam_w_mode, bias_correction1, bias_correction2,
+                         weight_decay):
         grads, params, exp_avgs, exp_avg_sqs = tensor_lists
+        for grad, param, exp_avg, exp_avg_sq in zip(grads, params, exp_avgs, exp_avg_sqs):
+            # The kernel indexes flat storage, so every operand must share the param's contiguous layout.
+            if not (param.is_contiguous() and grad.is_contiguous() and exp_avg.is_contiguous()
+                    and exp_avg_sq.is_contiguous()):
+                MPSFusedAdam._foreach_adam_step([[grad], [param], [exp_avg], [exp_avg_sq]], lr, beta1, beta2, epsilon,
+                                                adam_w_mode, bias_correction1, bias_correction2, weight_decay)
+                continue
+            kernel = MPSFusedAdam.kernels.get(param.dtype)
+            if kernel is None:
+                MPSFusedAdam._foreach_adam_step([[grad], [param], [exp_avg], [exp_avg_sq]], lr, beta1, beta2, epsilon,
+                                                adam_w_mode, bias_correction1, bias_correction2, weight_decay)
+                continue
+            kernel(param,
+                   grad,
+                   exp_avg,
+                   exp_avg_sq,
+                   float(lr),
+                   float(beta1),
+                   float(beta2),
+                   float(epsilon),
+                   float(weight_decay),
+                   float(bias_correction1),
+                   float(bias_correction2),
+                   float(adam_w_mode),
+                   threads=param.numel())
 
-        bias_correction1 = 1.0
-        bias_correction2 = 1.0
-        if bias_correction:
-            bias_correction1 = 1.0 - beta1**step
-            bias_correction2 = 1.0 - beta2**step
+    @staticmethod
+    def _foreach_adam_step(tensor_lists, lr, beta1, beta2, epsilon, adam_w_mode, bias_correction1, bias_correction2,
+                           weight_decay):
+        grads, params, exp_avgs, exp_avg_sqs = tensor_lists
+        if params[0].dtype != torch.float32:
+            # Match the kernel contract (fp32 math, storage-dtype results); fp16 intermediates overflow otherwise.
+            fp32_lists = [[t.float() for t in tensors] for tensors in tensor_lists]
+            MPSFusedAdam._foreach_adam_step(fp32_lists, lr, beta1, beta2, epsilon, adam_w_mode, bias_correction1,
+                                            bias_correction2, weight_decay)
+            for originals, fp32_tensors in zip((params, exp_avgs, exp_avg_sqs), fp32_lists[1:]):
+                torch._foreach_copy_(originals, fp32_tensors)
+            return
 
         # L2 mode folds weight decay into the gradient; AdamW mode applies it to the parameter directly.
         if weight_decay != 0 and not adam_w_mode:
@@ -52,7 +97,7 @@ def _adam_step(tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode, bias_
         torch._foreach_addcdiv_(params, exp_avgs, denom, value=-lr / bias_correction1)
 
 
-class FusedAdamBuilder(MPSOpBuilder):
+class FusedAdamBuilder(MetalOpBuilder):
     BUILD_VAR = "DS_BUILD_FUSED_ADAM"
     NAME = "fused_adam"
 
@@ -62,5 +107,30 @@ def __init__(self):
     def absolute_name(self):
         return f'deepspeed.ops.adam.{self.NAME}_op'
 
+    def metal_sources(self):
+        return ['csrc/mps/fused_adam.metal']
+
+    def is_compatible(self, verbose=False):
+        # The foreach fallback keeps FusedAdam usable on any MPS build.
+        return MPSOpBuilder.is_compatible(self, verbose)
+
     def load(self, verbose=True):
+        if MPSFusedAdam.kernels is None and not MPSFusedAdam.compile_failed and MetalOpBuilder.is_compatible(self):
+            try:
+                library = self.load_metal_library()
+            except RuntimeError as e:
+                # Keep the foreach path usable rather than failing the optimizer on a shader compile error.
+                MPSFusedAdam.compile_failed = True
+                self.warning(f"Metal FusedAdam kernel failed to compile, using torch._foreach fallback: {e}")
+                return MPSFusedAdam
+            MPSFusedAdam.kernels = {
+                torch.float32: library.fused_adam_float,
+                torch.float16: library.fused_adam_half,
+            }
+            # The bfloat kernel is only compiled on Metal 3.1+ (macOS 14); the library raises rather
+            # than returning an AttributeError for a missing entry point.
+            try:
+                MPSFusedAdam.kernels[torch.bfloat16] = library.fused_adam_bfloat
+            except RuntimeError:
+                pass
         return MPSFusedAdam