## Summary
This PR is the first step (phase 0) of enabling Apple Silicon support
for DeepSpeed: make single-device training work end to end with
pure-PyTorch ops.
[**To-Do in phase 1**] Metal kernels will come later and plug into the
`op_builder/mps` classes added here.
The MPS accelerator was a stub: memory queries returned `None`, no
communication backend was set, `fp16/bf16` were reported unsupported,
and every op builder resolved to `NotImplementedBuilder`.
`deepspeed.initialize` + one training step failed for every ZeRO stage
on an Apple Silicon machine. This PR aims on enabling capabilities.
### Changes
- **`accelerator/mps_accelerator.py`** — real `torch.mps` memory stats,
fp16/bf16 support (bf16 gated on macOS 14+), `torch.mps.Event`, `gloo`
as the comm backend, and `is_synchronized_device() = True` (PyTorch's
MPS backend effectively exposes a single in-order execution stream and
currently provides no public CUDA-style stream API or `record_stream`
mechanism.). Unified memory makes `pin_memory` a no-op (torch's
`pin_memory()` also raises under MPS).
- **`deepspeed/comm/torch.py`** — gloo cannot operate on MPS tensors
(even at world size 1), so collectives stage MPS tensors through CPU
copies via a `stage_on_cpu` decorator; async ops copy back on `wait()`.
- **`accelerator/abstract_accelerator.py` + `runtime/zero`** — MPS has
no fp64. Gradient-norm accumulation now picks its dtype via a new
concrete `is_fp64_supported()` (default `True`) and `get_norm_dtype()`
instead of hard-coded `.double()`.
- **`op_builder/mps/`** — new backend package (`MPSOpBuilder`,
`NotImplementedBuilder`, `FusedAdamBuilder`). `FusedAdam` is implemented
with `torch._foreach_*` ops and mirrors the math in
`csrc/adam/multi_tensor_adam.cu`, following the HPU precedent of
Python-backed builders.
- **`tests/unit/common.py`** — MPS must use `spawn` (Metal's compiler
service is lost in `forkserver` children, which hangs the harness) and
reports its device count via the accelerator.
- **`tests/unit/ops/adam/test_adamw.py`** —
`test_fused_adam_matches_torch` checks `FusedAdam` against
`torch.optim.Adam/AdamW` on the active accelerator (fp32/bf16 ×
Adam/AdamW), so it also guards the CUDA kernel.
### Verified on an M5 Max (macOS 26.3, torch 2.13.0)
- ZeRO 1/2/3 × fp32/bf16/fp16 train end to end with
`deepspeed.initialize` (single process). Also tested on ZeRO stage 0
which disables ZeRO completely and falling back to standard data
parallelism.
- MPS `FusedAdam` matches `torch.optim` to 2e-7 in fp32.
- `DS_ACCELERATOR=mps pytest unit/runtime/test_ds_config_dict.py
unit/runtime/test_ds_initialize.py
unit/runtime/half_precision/test_fp16.py
unit/runtime/half_precision/test_dynamic_loss_scale.py
unit/runtime/zero/test_zero_grad_clip.py
unit/runtime/zero/test_zero_context.py
unit/checkpoint/test_zero_optimizer.py`: 132 passed, 0 failed, 126
skipped (multi-device tests; `device_count() == 1`).
### Known limitations / follow-ups
- bf16 `FusedAdam` differs from the CUDA kernel by ~1 bf16 ulp (CUDA
computes in fp32 and stores bf16; the `_foreach` path rounds in bf16).
- The CPU-staged gloo path is only exercised at world size 1 here;
multi-Mac runs are untested.
- Follow-ups: macOS arm64 CI workflow, arm64 build of CPU Adam for
ZeRO-Offload, Metal kernels via `torch.mps.compile_shader`, and an Apple
Silicon tutorial page.
---------
Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
Zhipeng Wang
64fcec6ba7