[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#8300delock merged 3 commits into
Conversation
Sorry, something went wrong.
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>
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 GitHubYour 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".
Sorry, something went wrong.
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>
| # 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?
Sorry, something went wrong.
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.
Sorry, something went wrong.
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.
Sorry, something went wrong.
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.
Sorry, something went wrong.
| 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?
Sorry, something went wrong.
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).
Sorry, something went wrong.
| 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.
Sorry, something went wrong.
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.
Sorry, something went wrong.
| 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?
Sorry, something went wrong.
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.
Sorry, something went wrong.
…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>
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
Verified on an M5 Max (macOS 26.3, torch 2.13.0)
Follow-ups