Conversation
The CPU Adam kernel ran scalar on AArch64 machines without SVE, which includes every Apple Silicon Mac (the ZeRO-Offload optimizer path on that platform). simd.h gains a 4-lane NEON implementation of the existing SIMD macro set, following the AVX512 semantics: fp16 through the hardware converters, bf16 through the same round-to-nearest-even with nan-quieting flow as store_16_f32_as_bf16_nearest, and x86 andnot's operand order preserved. Step_AVX's non-AVX512 bf16 bailout is lifted for NEON since the bf16 converters exist there. Builders: simd_width() advertises -D__NEON__ for AArch64 without SVE (the vdivq/vsqrtq intrinsics used are A64-only, so 32-bit ARM keeps the scalar path) and the MPS cpu_adam builder switches from -D__SCALAR__ accordingly. Lion/Adagrad gates are left untouched and keep their scalar behavior. Measured on an M5 Max (macOS 26.3, clang, libomp), DeepSpeedCPUAdam step on 50M params, 10-step average, vs the -D__SCALAR__ build: fp32 11.5 -> 3.8 ms, fp16 11.5 -> 3.3 ms, bf16 12.7 -> 4.9 ms. NEON and scalar builds produce bit-identical fp16 results on identical inputs; all dtypes match a fp32 torch.optim oracle within storage rounding, including sizes exercising the SIMD/scalar tail boundary. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: dfd95919b5
ℹ️ 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".
| { | ||
| size_t rounded_size = 0; | ||
| #if defined(__AVX512__) or defined(__AVX256__) | ||
| #if defined(__AVX512__) or defined(__AVX256__) or defined(__NEON__) |
There was a problem hiding this comment.
Use __NEON__ to filter may have impact on SVE device, @xylian86 to confirm.
There was a problem hiding this comment.
@delock the two defines can't coexist: they're both emitted only by simd_width(), which checks 'sve' in cpuinfo flags first and returns -D__SVE__ there; -D__NEON__ is only returned for ARM_8 without sve, so an SVE device never sees the NEON define. In the source, every '#if ... or defined(NEON)' gate sits before the existing '#elif defined(SVE)' branch, and with NEON undefined those reduce to exactly the pre-PR expressions — the SVE dispatch, its Step_SVE path, and its log line are untouched. Happy to adjust if @xylian86 sees an SVE configuration I've missed.
cxx_args hardcoded -std=c++17, which lands after cpp_extension's own -std flag and overrides it. torch 2.14's headers require C++20, so the JIT build broke on CI runners with current torch while still compiling against torch 2.13 locally. Drop the flag; cpp_extension supplies the standard its headers need. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
|
Hi @PKUWZP , do you think using some different macro as NEON as path selection? NEON is built-in macro when running ARM compiler. I'm worrying that for ARM chips with SVE feature, the code here would direct C++ code into neon path which may not be intended. Use a macro slightly different would be a better choice. |
create_adam_optimizer's diagnostic reported scalar arithmetic on NEON builds; add the missing branch. The ZeRO-Offload tuning guide now describes the NEON path for AArch64 CPUs without SVE instead of claiming they stay scalar. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
Resolve op_builder/mps/cpu_adam.py: keep master's torch-version-aware C++ standard selection (its fix for the same torch 2.14 header issue) with this branch's -D__NEON__. Signed-off-by: PKUWZP <zhipeng.rainbowserie@gmail.com>
Summary
Phase 2 of Apple Silicon support (follow-up to #8293/#8300/#8335): the CPU Adam kernel — the ZeRO-Offload optimizer path — ran scalar on AArch64 machines without SVE, which includes every Apple Silicon Mac. This adds a 4-lane NEON implementation of the existing SIMD macro layer.
Changes
csrc/includes/simd.h— a__NEON__branch defining the full macro set (SIMD_LOAD/STORE/SET/ADD/MUL/FMA/SQRT/DIV/AND/ANDNOT/OR/XOR, width 4):vcvt_f32_f16/vcvt_f16_f32).store_16_f32_as_bf16_nearest(usingvaddhn_u32for the add-and-take-high-half step); loads are widen+shift.andnot(x, y) = ~x & ymaps tovbicq(y, x)— operand order preserved (documented in a comment).simd_load/simd_storeguards widen from AVX512-only to AVX512-or-NEON.csrc/includes/cpu_adam.h—Step_AVX's non-AVX512 bf16 bailout is lifted for NEON (this was silently sending bf16 back to the scalar tail); the two Adam gates widen to include__NEON__.csrc/adam/cpu_adam_impl.cpp— same gate widening (4 sites, includingkZenAdamAlign). The NEON branch sits before the existing__SVE__alternative and they remain mutually exclusive builder-emitted defines.op_builder/builder.py—simd_width()advertises-D__NEON__forARM_8without SVE. 32-bit ARM keeps__SCALAR__: thevdivq_f32/vsqrtq_f32intrinsics used are A64-only.op_builder/mps/cpu_adam.py— switches from-D__SCALAR__to-D__NEON__.Measured on Apple M5 Max (macOS 26.3, Apple clang, Homebrew libomp)
DeepSpeedCPUAdamstep, 50M params, 10-step average, vs the-D__SCALAR__build of the same tree:Correctness
torch.optim.Adam/AdamWoracle within storage rounding, at sizes exercising pure-SIMD, SIMD+scalar-tail (1000003), and sub-width (3) paths.test_cpu_adam.py+test_hybrid_adam.py+test_adamw.py— 122 passed, 7 skipped. Themps-torch-latestCI workflow JIT-builds this kernel in its offload configs, so the NEON path is exercised upstream on every touching PR.