DeepSpeed

715965e0 - [Apple Silicon Support Phase 1] Add Metal FusedAdam kernel and CPU Adam build for Apple Silicon (#8300)

[Apple Silicon Support Phase 1] Add Metal FusedAdam kernel and CPU Adam build for Apple Silicon (#8300) ## Summary This PR implements the Phase 1 work of Apple Silicon support for DeepSpeed (follow-up to #8293, which made single-device ZeRO 1–3 training work with pure-PyTorch ops). This PR adds the **op-builder layer**: - Metal kernels compiled at runtime, - a first real Metal kernel (FusedAdam), - and the C++ CPU Adam build so ZeRO-Offload works on Macs. ### Changes - **`op_builder/mps/builder.py`** — new `MetalOpBuilder`. Subclasses list `.metal` files in `metal_sources()`; the shader is compiled at `load()` through `torch.mps.compile_shader`, which dispatches kernels on PyTorch's own MPS command stream. No Xcode project, `.metallib` packaging, or C++ extension build is involved. `is_compatible()` additionally requires `torch.mps.compile_shader`. - **`csrc/mps/fused_adam.metal` + `op_builder/mps/fused_adam.py`** — `FusedAdam` becomes a Metal kernel (one launch per tensor). It does all math in fp32 and stores in the parameter dtype, the same contract as `csrc/adam/multi_tensor_adam.cu`; this also closes the bf16 ulp gap noted in #8293. The `torch._foreach_*` implementation remains as a fallback for torch builds without `compile_shader` and for non-contiguous tensors. - **`op_builder/mps/cpu_adam.py`** — builds `csrc/adam/cpu_adam*.cpp` with the system clang (`-D__SCALAR__` on arm64). Apple clang has no OpenMP, so the build uses Homebrew `libomp` when `brew --prefix libomp` resolves and omits it otherwise; both paths verified. This enables `DeepSpeedCPUAdam` and therefore ZeRO-Offload on Apple Silicon. Unified memory means offloading does not copy parameters between separate memories. - **`tests/unit/ops/adam/test_adamw.py`** — `test_fused_adam_matches_reference` checks `FusedAdam` against an explicit fp32-math / storage-dtype-rounding reference for fp32, bf16, and fp16 × Adam/AdamW (replaces the `torch.optim` comparison from #8293, whose bf16 reference computes in bf16 and is a worse baseline). Tolerance is 8 ulp of the storage dtype at tensor scale, which covers measured fp32 op-order drift over 5 steps. - **`tests/unit/ops/adam/test_cpu_adam.py`, `test_hybrid_adam.py`** — `py-cpuinfo` has no `vendor_id_raw` on Apple Silicon; use `.get()`. - **`MANIFEST.in`** — ship `.metal` sources. **Docs** — accelerator setup guide updated for offload and the Metal/OpenMP notes. ### Verified on an M5 Max (macOS 26.3, torch 2.13.0) - `DeepSpeedCPUAdam` matches `torch.optim` to ~1e-6; ZeRO-Offload trains end to end for stage 1/2/3 × fp32/bf16/fp16 (optimizer offload; plus param offload for stage 3). - Metal `FusedAdam` vs foreach fallback, 4×100k params: 0.08 vs 0.27 ms/step (fp32), 0.03 vs 0.16 ms/step (bf16). Both implementations pass the new reference test in all 6 cases. - `DS_ACCELERATOR=mps pytest unit/ops/adam/test_cpu_adam.py unit/ops/adam/test_hybrid_adam.py unit/ops/adam/test_adamw.py`: 74 passed, 7 skipped. - `op_builder.mps` imports with torch absent (the sdist/install-smoke path). ### Follow-ups - macOS arm64 CI workflow so these paths are exercised upstream. - Further Metal kernels (quantizer for ZeRO++). --------- Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>