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. |
PKUWZP
changed the title
![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: 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.
This non-merge commit has no Signed-off-by trailer, so it violates the repository's commit requirement and will be rejected by the DCO check; recreate the commit using --signoff.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
| elif cpu_info['arch'] == 'ARM_8': | ||
| # NEON is baseline on AArch64; vdivq/vsqrtq used by simd.h are A64-only, | ||
| # so do not advertise it for 32-bit ARM. | ||
| return '-D__NEON__' |
There was a problem hiding this comment.
Returning __NEON__ here makes non-SVE AArch64 builds vectorized for fp32, fp16, and bf16, but docs/_tutorials/zero-offload.md:78 still explicitly tells users that these CPUs and non-fp32 types use the scalar implementation. Update that guidance as part of this feature so users do not make tuning decisions from the now-incorrect capability description.
AGENTS.md reference: AGENTS.md:L26-L26
Useful? React with 👍 / 👎.
| { | ||
| 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.
When info logging is enabled, this new __NEON__ dispatch executes the vector kernel, but create_adam_optimizer has no corresponding __NEON__ logging branch and therefore prints that the optimizer was created with scalar arithmetic. This makes the normal DeepSpeedCPUAdam startup diagnostic incorrect on every affected AArch64 build; add a NEON case alongside the AVX and SVE cases.
Useful? React with 👍 / 👎.
| { | ||
| 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.
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. |
2 participants
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.