## 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>
With Intel Extension for PyTorch retiring, XPU device would be supported
by PyTorch 2.8+ and dependency to Intel Extension for PyTorch would not
be needed.
This PR removed IPEX dependency, adapt to builder protocol in PyTorch
for XPU, and updated documents and tests accordingly.
Note after this update, DeepSpeed will not work with previous
PyTorch+IPEX on XPU devices. Suggest user to upgrade to latest PyTorch
to get latest XPU features on XPU devices.
Come with this PR is removal of InferenceBuilder, the kernel needed by
InferenceBuilder is supported through Intel Extension for PyTorch.
---------
Signed-off-by: Ma, Guokai <guokai.ma@intel.com>
Co-authored-by: Olatunji Ruwase <tjruwase@gmail.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Two spelling errors in docs/_tutorials/accelerator-setup-guide.md:
Line 50: comma-seperated-dash-range → comma-separated-dash-range
Line 97: optimzied → optimized
Both typos are in the Intel Architecture CPU section of the accelerator
setup guide.
Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com>
This document provide a places to hold accelerator setup guides. It is
intend to be a single place to lookup installation guides for different
accelerators. Currently CPU and XPU setup guides are added to this
document and could be extended to other accelerators.
---------
Co-authored-by: Liangliang Ma <1906710196@qq.com>
Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
Zhipeng Wang