[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,55 @@
+// 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);
+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);

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,60 @@
 
 
 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
 
     @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[param.dtype]
+            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
 
         # 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 +84,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 +94,19 @@ 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 MetalOpBuilder.is_compatible(self):
+            library = self.load_metal_library()
+            MPSFusedAdam.kernels = {
+                torch.float32: library.fused_adam_float,
+                torch.float16: library.fused_adam_half,
+                torch.bfloat16: library.fused_adam_bfloat,
+            }
         return MPSFusedAdam

tests/unit/ops/adam/test_adamw.py

@@ -80,32 +80,50 @@ def test(self,
             assert ds_optimizer.adam_w_mode == adam_w_mode
 
 
+def reference_adam_step(param, grad, exp_avg, exp_avg_sq, step, lr, beta1, beta2, eps, weight_decay, adam_w_mode):
+    """Adam/AdamW step with fp32 math and storage-dtype rounding, matching csrc/adam/multi_tensor_adam.cu."""
+    dtype = param.dtype
+    p, g, m, v = param.float(), grad.float(), exp_avg.float(), exp_avg_sq.float()
+    if not adam_w_mode:
+        g = g + weight_decay * p
+    m = beta1 * m + (1 - beta1) * g
+    v = beta2 * v + (1 - beta2) * g * g
+    denom = (v / (1 - beta2**step)).sqrt() + eps
+    update = (m / (1 - beta1**step)) / denom
+    if adam_w_mode:
+        update = update + weight_decay * p
+    p = p - lr * update
+    param.copy_(p.to(dtype))
+    exp_avg.copy_(m.to(dtype))
+    exp_avg_sq.copy_(v.to(dtype))
+
+
 @pytest.mark.parametrize('adam_w_mode', [True, False], ids=["adamw", "adam"])
-@pytest.mark.parametrize('dtype', [torch.float, torch.bfloat16], ids=["fp32", "bf16"])
-def test_fused_adam_matches_torch(adam_w_mode, dtype):
+@pytest.mark.parametrize('dtype', [torch.float, torch.bfloat16, torch.half], ids=["fp32", "bf16", "fp16"])
+def test_fused_adam_matches_reference(adam_w_mode, dtype):
     if dtype not in get_accelerator().supported_dtypes():
         pytest.skip(f"{dtype} not supported on {get_accelerator().device_name()}")
     if not deepspeed.ops.__compatible_ops__[FusedAdamBuilder.NAME]:
         pytest.skip("FusedAdam is not compatible")
 
     device = get_accelerator().device_name()
     torch.manual_seed(0)
-    ref_params = [torch.randn(1024, device=device, dtype=dtype, requires_grad=True) for _ in range(3)]
-    ds_params = [torch.nn.Parameter(p.detach().clone()) for p in ref_params]
-    optimizer_kwargs = dict(lr=1e-2, weight_decay=0.1)
-    torch_optimizer = torch.optim.AdamW if adam_w_mode else torch.optim.Adam
-    ref_optimizer = torch_optimizer(ref_params, **optimizer_kwargs)
-    ds_optimizer = FusedAdam(ds_params, adam_w_mode=adam_w_mode, **optimizer_kwargs)
+    lr, betas, eps, weight_decay = 1e-2, (0.9, 0.999), 1e-8, 0.1
+    ds_params = [torch.nn.Parameter(torch.randn(1024, device=device, dtype=dtype)) for _ in range(3)]
+    ref_params = [p.detach().clone() for p in ds_params]
+    ref_exp_avgs = [torch.zeros_like(p) for p in ref_params]
+    ref_exp_avg_sqs = [torch.zeros_like(p) for p in ref_params]
+    ds_optimizer = FusedAdam(ds_params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, adam_w_mode=adam_w_mode)
 
-    for _ in range(5):
-        for ref_param, ds_param in zip(ref_params, ds_params):
-            grad = torch.randn_like(ref_param)
-            ref_param.grad = grad.clone()
-            ds_param.grad = grad.clone()
-        ref_optimizer.step()
+    for step in range(1, 6):
+        for ds_param, ref_param, exp_avg, exp_avg_sq in zip(ds_params, ref_params, ref_exp_avgs, ref_exp_avg_sqs):
+            ds_param.grad = torch.randn_like(ds_param)
+            reference_adam_step(ref_param, ds_param.grad, exp_avg, exp_avg_sq, step, lr, betas[0], betas[1], eps,
+                                weight_decay, adam_w_mode)
         ds_optimizer.step()
 
-    # bf16 storage rounds differently depending on where intermediates are kept, so allow one ulp.
-    atol = 1e-5 if dtype == torch.float else 2e-2
-    for ref_param, ds_param in zip(ref_params, ds_params):
-        torch.testing.assert_close(ds_param.float(), ref_param.float(), atol=atol, rtol=0)
+    # fp32 operation order differs between implementations and accumulates over steps, so allow a
+    # few ulps of the storage dtype at the scale of the tensor (per-element rtol is too strict near zero).
+    for ds_param, ref_param in zip(ds_params, ref_params):
+        atol = 8 * torch.finfo(dtype).eps * ref_param.abs().max().item()
+        torch.testing.assert_close(ds_param.float(), ref_param.float(), rtol=0, atol=atol)

tests/unit/ops/adam/test_cpu_adam.py

@@ -19,7 +19,7 @@
 if not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]:
     pytest.skip("cpu-adam is not compatible", allow_module_level=True)
 
-pytest.cpu_vendor = get_cpu_info()["vendor_id_raw"].lower()
+pytest.cpu_vendor = get_cpu_info().get("vendor_id_raw", "").lower()
 
 
 def check_equal(first, second, atol=1e-2, verbose=False):