DeepSpeed
11ac9f86 - Add RVV vectorized FP32 kernel to CPU Adam (#8817)

Commit
14 days ago
Add RVV vectorized FP32 kernel to CPU Adam (#8817) ## Description On riscv64 hosts with the RISC-V Vector (RVV 1.0) extension, DeepSpeed CPU Adam currently falls back to the scalar `Step_1` tail because neither the op builder's `simd_width()` nor the kernel dispatch recognizes the V extension. This PR adds an RVV FP32 tier to CPU Adam, mirroring the existing SVE backend, so RVV-capable CPUs get a ~4x single-core speedup on the optimizer step. - `csrc/includes/rvv.h`: an RVV 1.0 FP32 Adam kernel. The hot loop processes one `m4` block per iteration — at VLEN=128 each block covers a full 64-byte cache line per stream — and runs entirely at the loop-invariant vlmax; a single `vsetvl` truncates the trailing block, so any tensor length is covered with no scalar tail. The arithmetic order matches the scalar and SVE kernels, keeping every tier within the existing test tolerance. Scalar hyperparameters ride the vf-form intrinsics, so only epsilon occupies a vector register. - `cpu_adam.h` / `cpu_adam_impl.cpp`: `Step_RVV` dispatch alongside `Step_AVX` and `Step_SVE`; only the FP32 param/FP32 state pair is vectorized, other dtype pairs keep the scalar path (same policy as SVE). - `op_builder/builder.py`: detect the V extension from the `/proc/cpuinfo` isa line with a per-character token test (a substring test misreads `rv64imafdcvh` as having V), build with `-D__RVV__ -march=rv64gcv` on riscv64, and normalize py-cpuinfo's `RISCV_64` arch string so the riscv64 branch actually fires. Hosts without V are unaffected (`rv64gc`). - `docs/_tutorials/zero-offload.md`: document the automatic RVV tier. ## Validation **Hardware/software:** SG2044 60-core RISC-V server, VLEN=128, gcc-14.3, torch 2.15 CPU-only build. **Correctness:** `pytest tests/unit/ops/adam/test_cpu_adam.py` → **48 passed, 1 skipped, 0 failed**. The suite compares CPU Adam against the PyTorch fused-Adam reference across fp32/fp16 and multiple sizes; the fp16 equality tests pass (this tier keeps non-fp32 dtype pairs on the scalar path). Key log: ``` 48 passed, 1 skipped, 3 warnings in 1283.18s (0:21:23) ``` **Performance** (single core, `OMP_NUM_THREADS=1 taskset -c 0`, 100 steps): | Size | Scalar baseline | RVV (this PR) | Speedup | |------|-----------------|---------------|---------| | 1B params | 13.07 s/step | 3.15 s/step | ~4.2x | | 2B params | 26.13 s/step | 6.31 s/step | ~4.1x | Key log: ``` Elapsed Time is 3.1488167250156405 Elapsed Time is 6.313579152822495 ``` --------- Signed-off-by: 张健10355098 <zhang.jian60@zte.com.cn>
Author
Parents
Loading