[Apple Silicon Support Phase 1] Add Metal FusedAdam kernel and CPU Adam build for Apple Silicon - #8300
Conversation
Phase 1 of Apple Silicon support: the op builder layer. - op_builder/mps/builder.py: MetalOpBuilder compiles .metal sources at load time through torch.mps.compile_shader, so Metal kernels dispatch on PyTorch's MPS stream without an Xcode project or C++ extension. - csrc/mps/fused_adam.metal + op_builder/mps/fused_adam.py: FusedAdam becomes a Metal kernel that does its math in fp32 and stores in the parameter dtype, matching csrc/adam/multi_tensor_adam.cu. It is 3-5x faster than the torch._foreach path, which remains as the fallback for torch builds without compile_shader. - op_builder/mps/cpu_adam.py: build the C++ CPU Adam kernel with clang so ZeRO-Offload works on Macs. OpenMP is enabled when Homebrew libomp is present and silently omitted otherwise. - tests/unit/ops/adam/test_adamw.py: check FusedAdam against an fp32 reference with storage-dtype rounding, for fp32/bf16/fp16. - tests/unit/ops/adam: py-cpuinfo has no vendor_id_raw on Apple Silicon. - MANIFEST.in ships .metal sources; docs updated. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a21346d887
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
PKUWZP
changed the title
bfloat is a Metal 3.1 (macOS 14) type, so compiling the whole shader eagerly failed on older systems and took the float/half kernels down with it. Gate the bfloat specialization on __METAL_VERSION__, pick up its entry point only when present, and fall back to the foreach path (with a single warning) if the shader fails to compile at all. The foreach fallback now does its math in fp32 for half-precision parameters, matching the kernel contract; fp16 intermediates overflowed. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
PKUWZP
changed the title
| # 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() |
There was a problem hiding this comment.
Is this drift impact fp32 only? Should we retain original atol for bf16?
There was a problem hiding this comment.
What is the new atol value compared to old atol value (1e-5)? An example might help to see how much is relaxed here.
There was a problem hiding this comment.
The fp32 op-order drift affects every dtype (both implementations do their math in fp32), but for bf16/fp16 the dominant term in any honest bound is the storage rounding: one boundary flip moves an element by a full storage ulp, which dwarfs the drift. The old 2e-2 was calibrated against a different reference (torch.optim running bf16 math) that this PR replaces, so it isn't directly comparable — measured agreement against the new fp32-math reference is ~3e-5 for bf16, far inside the bound. I've added the concrete numbers to the comment in 34f7cb7; happy to tighten bf16/fp16 to fewer ulps if you'd prefer a snugger bound.
There was a problem hiding this comment.
Good idea — added to the code comment in 34f7cb7. Concretely, for this test's data (|param| ~ 3 after 5 steps): fp32 atol evaluates to ~3e-6 (tighter than the old 1e-5), bf16 to ~0.2, fp16 to ~2.5e-2, while the measured implementation-vs-reference differences are ~1e-6 (fp32) and ~3e-5 (bf16). So fp32 got stricter, and the loose-looking bf16 bound is headroom over a much smaller observed error.
| 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; } |
There was a problem hiding this comment.
why adam_w_mode is a float?
There was a problem hiding this comment.
No good reason — at the time I hadn't verified that the shader binding marshals Python ints. It does: fixed in 34f7cb7, adam_w_mode is now constant uint& and the Python side passes int(adam_w_mode).
| 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. |
There was a problem hiding this comment.
One question is for unified memory whether offloading to CPU optimizer is necessary.
There was a problem hiding this comment.
Fair question — on unified memory offload is not about capacity at all (optimizer states occupy the same DRAM either way). What it still buys: Metal caps a process's GPU working set below total RAM (torch.mps.recommended_max_memory, ~75% here), and CPU-held optimizer state stays outside that budget; the step also runs on the CPU cores. So: unnecessary when the model fits the working-set budget, useful when that limit binds. Reworded the doc to say exactly this in 34f7cb7.
| return hasattr(torch.backends, "mps") and torch.backends.mps.is_available() | ||
|
|
||
|
|
||
| class MetalOpBuilder(MPSOpBuilder): |
There was a problem hiding this comment.
is that true that ops on Metal are not compiled and saved on disk cache? Is it temporary or its nature of Metal?
There was a problem hiding this comment.
It's the nature of the Metal workflow rather than a temporary gap: torch.mps.compile_shader compiles these small kernels in milliseconds at first load, and macOS's Metal framework maintains its own per-app on-disk cache of compiled pipelines, so a torch-extensions style build cache would add complexity without saving anything. Documented in the MetalOpBuilder docstring in 34f7cb7.
…notes - adam_w_mode reaches the Metal kernel as constant uint& instead of a float flag; the shader binding marshals Python ints natively. - Spell out what the ulp-scaled test tolerance evaluates to against the measured implementation agreement. - Docs: on unified memory offload is not about capacity; it moves the optimizer state out of Metal's GPU working-set budget and the step onto the CPU cores. - MetalOpBuilder: note why there is no torch-extensions disk cache (millisecond runtime compiles; Metal keeps its own pipeline cache). Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
Resolve tests/unit/ops/adam/test_adamw.py in favor of master's reference-based FusedAdam test from #8300, which supersedes the bf16 trim this branch carried for the old torch-comparison test. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
) ## Summary The CPU `fused_adam` extension created its `Adam_Optimizer` once with default arguments and ignored the `mode` parameter entirely (`csrc/cpu/adam/fused_adam.cpp`), so `FusedAdam(adam_w_mode=False)` on the CPU backend always applied decoupled (AdamW) weight decay instead of L2. Everything else (`lr`, betas, `eps`, `weight_decay`, bias correction) is already passed per call via `ds_adam_step`; only the AdamW-vs-L2 flag is fixed at construction. The fix keeps one optimizer instance per mode (`mode 1 == AdamW`, matching the CUDA kernel's `ADAM_MODE_1`). Also trims `test_fused_adam_matches_torch` to fp32: its bf16 cases compared against `torch.optim` running bf16 math, while the fused kernels compute in fp32 — never a valid reference. Low-precision dtypes get an explicit fp32-math reference test in the FusedAdam rework (deepspeedai#8300). ## How this surfaced Split out of deepspeedai#8303 at @delock's request: after a master merge, cpu-torch-latest failed on `test_fused_adam_matches_torch[fp32-adam]` (98.7% of elements mismatched — systematic, not tolerance noise), and the investigation traced it to this binding. The fix was verified green on cpu-torch-latest in deepspeedai#8303's CI (run 32695...) before being extracted here. ## Validation - `test_fused_adam_matches_torch[fp32-adam]` / `[fp32-adamw]` now genuinely exercise both decay modes against `torch.optim.Adam` / `AdamW` on the active accelerator. - cpu-torch-latest passed with this exact change as part of deepspeedai#8303's branch; this PR carries it alone. --------- Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
…erator (deepspeedai#8335) ## Summary Closes the remaining gap in the Apple Silicon support series (deepspeedai#8293, deepspeedai#8300, deepspeedai#8303, deepspeedai#8307): none of the MPS paths were exercised by CI — every MPS-gated test skips on Linux runners, so regressions could only be caught on a developer's Mac. ### macOS CI workflow (`mps-torch-latest.yml`) Runs the MPS-green unit test subset on GitHub's arm64 macOS runners (`macos-15`), which expose a working MPS device: - `unit/ops/adam/test_adamw.py` — Metal/foreach FusedAdam vs fp32-math reference, CPU Adam configs incl. ZeRO-Offload - `unit/comm/test_dist.py` — gloo CPU-staging for collectives and P2P (`TestMpsStagedP2P`) - `unit/runtime/test_ds_config_dict.py` — config-driven `deepspeed.initialize` + training steps **Designed not to interfere with existing CI:** - PR triggers are scoped via `paths:` to MPS-relevant files (`accelerator/**`, `op_builder/mps/**`, `csrc/mps/**`, `deepspeed/comm/**`, the two test dirs, and the workflow itself) — the check does not even appear on unrelated PRs. - Separate workflow, own concurrency group with cancel-in-progress, hard `timeout-minutes: 45`. - Not a required check (that's a branch-protection setting; nothing here changes it), so even a red run cannot block merges of non-macOS work. - Nightly `schedule` + `workflow_dispatch` for coverage between touching PRs. ### torch floor check `MPS_Accelerator.__init__` now fails with a clear message on torch older than 2.3, where the `torch.mps` memory queries ZeRO depends on (`recommended_max_memory`) do not exist — previously this surfaced as a bare `AttributeError` deep inside ZeRO's flatten logic. Feature-detected rather than version-parsed. (The Metal FusedAdam kernel already degrades gracefully on torch without `compile_shader`.) ## Validation - The workflow's exact pytest command passes locally on an M5 Max (macOS 26.3, torch 2.13): 65 passed, 23 skipped (multi-device), 1m54s — comfortably inside the runner budget. - Guard verified both ways: normal construction unaffected; with `recommended_max_memory` hidden, construction raises the explicit `ValueError`. --------- Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
2 participants
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:
Changes
op_builder/mps/builder.py— newMetalOpBuilder. Subclasses list.metalfiles inmetal_sources(); the shader is compiled atload()throughtorch.mps.compile_shader, which dispatches kernels on PyTorch's own MPS command stream. No Xcode project,.metallibpackaging, or C++ extension build is involved.is_compatible()additionally requirestorch.mps.compile_shader.csrc/mps/fused_adam.metal+op_builder/mps/fused_adam.py—FusedAdambecomes a Metal kernel (one launch per tensor). It does all math in fp32 and stores in the parameter dtype, the same contract ascsrc/adam/multi_tensor_adam.cu; this also closes the bf16 ulp gap noted in Enable DeepSpeed support on Apple Silicon (MPS) with ZeRO Stage 1-3 #8293. Thetorch._foreach_*implementation remains as a fallback for torch builds withoutcompile_shaderand for non-contiguous tensors.op_builder/mps/cpu_adam.py— buildscsrc/adam/cpu_adam*.cppwith the system clang (-D__SCALAR__on arm64). Apple clang has no OpenMP, so the build uses Homebrewlibompwhenbrew --prefix libompresolves and omits it otherwise; both paths verified. This enablesDeepSpeedCPUAdamand 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_referencechecksFusedAdamagainst an explicit fp32-math / storage-dtype-rounding reference for fp32, bf16, and fp16 × Adam/AdamW (replaces thetorch.optimcomparison from Enable DeepSpeed support on Apple Silicon (MPS) with ZeRO Stage 1-3 #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-cpuinfohas novendor_id_rawon Apple Silicon; use.get().MANIFEST.in— ship.metalsources. Docs — accelerator setup guide updated for offload and the Metal/OpenMP notes.Verified on an M5 Max (macOS 26.3, torch 2.13.0)
DeepSpeedCPUAdammatchestorch.optimto ~1e-6; ZeRO-Offload trains end to end for stage 1/2/3 × fp32/bf16/fp16 (optimizer offload; plus param offload for stage 3).FusedAdamvs 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.mpsimports with torch absent (the sdist/install-smoke path).Follow-ups