Peng add cpu adam fused step v1

[pengdurice]

ZeRO-3 CPU offload: apply the clip factor and write the bf16 copy inside the CPU Adam pass

Builds on: #TODO(pinned copy-back PR number) (ZeRO-3: copy offloaded master weights back through pinned buffers).
This branch contains that PR's commit followed by this one; review the second commit.
Part of: #8672 (Q4 2026 roadmap), "Enable training of large scale models".

Summary

With the optimizer offloaded to the CPU, ZeRO-3's step makes three passes over every CPU sub-group:

1. multiply every fp32 gradient by the clipping factor in place (grad.mul_), even when the factor is 1;
2. run DeepSpeedCPUAdam (reads gradients, parameters and both moments; writes parameters and both moments);
3. cast the updated fp32 parameters to bf16 into the pinned staging buffers and copy them to the device.

This PR folds 1 and 3 into 2: the CPU Adam kernel multiplies each gradient by the factor as it loads it, and writes
the bf16 copy of each updated piece into a pinned staging buffer, which is copied to the device while the next piece
is updated. The parameters after training are bitwise equal to the three-pass path.

Duplicate check

Searched open and recent pull requests on 2026-10-15 ("cpu adam", "offload", "pinned", "master weights"): no PR
fuses the clip factor or the copy-back into the CPU Adam step. Related but different: #8673 fuses mixed-precision
gradients and the parameter writeback into the GPU FusedAdam; #8813 builds the CPU optimizer ops without CUDA and
fixes the CPU Lion update; #8817 adds an RVV kernel to CPU Adam (both change the same C++ files, so whichever merges
later rebases); #8719 adds asynchronous CPU-offloaded optimizers.

The change

File	Lines	What
csrc/includes/cpu_adam.h	+72 -1	_grad_scale, _params_copy, set_grad_scale, set_params_copy, params_copy_for: the step multiplies each loaded gradient by the factor and, when a copy buffer is set, stores the updated parameter in bf16 there
csrc/includes/sve.h	+4 -1	the SVE path takes the factor
csrc/adam/cpu_adam_impl.cpp	+117	ds_adam_step_scaled, ds_adam_step_scaled_copy, ds_adam_copy_supported, ds_adam_block_numel
csrc/adam/cpu_adam.cpp	+14	bindings adam_update_scaled, adam_update_scaled_copy, adam_copy_supported, adam_block_numel
deepspeed/ops/adam/cpu_adam.py	+66 -4	step(closure=None, grad_scale=1.0); copy_supported(); step_pieces(piece_numel, before_piece, after_piece, grad_scale=1.0) updates a parameter piece by piece, before_piece returns the bf16 buffer for the piece, after_piece starts its copy; the piece size is rounded down to a multiple of the kernel's widest vector block (adam_block_numel, 128 floats with AVX512) so every element takes the vector path, as in one call on the whole tensor
deepspeed/runtime/zero/stage3.py	+97 -10	_grad_scale_for_cpu_adam passes the factor to CPU Adam instead of mul_ for CPU sub-groups; _cpu_adam_step_and_copy_back runs step_pieces into the two pinned staging buffers of the copy-back PR, waits for the previous copy out of a buffer before refilling it, and waits for the last piece before returning (the next forward's all-gathers do not wait for this stream)
tests/unit/ops/adam/test_cpu_adam.py	+78	kernel tests below
tests/unit/v1/zero/test_zero.py	+118	ZeRO-3 tests below

Where the fused path runs: CPU sub-groups with float32 gradients and master weights, pin_memory on, a bf16 device
partition, and a CPU Adam build with a bf16 store (AVX512, NEON and the scalar loop; copy_supported() asks the build).
Everything else (fp16 or a GradScaler, ZenFlow, NVMe offload, Muon, SuperOffload, pin_memory off) keeps the current
path. No configuration key changes.

Measured

GLM-5.2, real weights, PG-19 rows of 65,536 tokens, fp32 Adam states, CPU optimizer offload, each sequence over 2 GPUs.

Setup	Without this PR	With this PR	Same losses
1 node x 8 H200, first 8 layers, 20% of the optimizer on the GPU: clip factor fused	15.95 s per step, optimizer step 5.01-5.29 s	15.01 s, optimizer step 4.25-4.33 s	yes
the same, clip factor and bf16 copy fused (2 runs)	15.95 s	14.61-14.67 s (-8.0 to -8.4%)	yes
8 nodes x 8 H200, all 78 layers, 10% on the GPU, one job on the same nodes	147.16 s	143.61 s (-2.4%)	yes (to the printed digits until step 8, then one bf16 step, as between repeated runs)

Profile (one step, rank 1): the separate aten::mul_ pass took 824 ms of a 6.70 s step on 2 nodes before;
with this PR it is 1.2 ms in an 8-node step, and no separate cast pass remains after the backward.
Memory: no new buffers (the staging buffers belong to the copy-back PR).

This change was also in the DeepSpeed tree of our MiniMax-M3 and GLM-5.2 runs against Megatron-LM (real weights,
65,536-token sequences, 80-95% of the optimizer on the CPU): MiniMax-M3 on 4 nodes 85.24 s per step vs 119.35 s for
Megatron-LM (one job on the same nodes); GLM-5.2 first 38 layers on 4 nodes 62.03 vs 63.73 s (one job on the same nodes);
GLM-5.2 all 78 layers on 8 nodes 130.58 vs 135.56 s (separate jobs on 7 of the same 8 nodes). On one node, master plus this
PR, the copy-back PR and the two model PRs (#8658, #8659) gives the same loss at every step as that tree.

Tests

Test	Checks
test_cpu_adam_grad_scale_matches_scaling_in_place	step(grad_scale=f) gives bitwise the parameters and moments of grad.mul_(f) then step()
test_cpu_adam_step_pieces_matches_step	step_pieces with pieces of 1,000 elements (not a multiple of the vector block) and 524,289-element tensors gives bitwise the result of one step(), and the bf16 copies equal param.to(torch.bfloat16)
TestZero3CpuAdamGradScale (2 ranks; clip 1e-3 and 1e6; ratio 1.0 and 0.5)	ZeRO-3 with CPU offload: parameters after 3 steps bitwise equal to the in-place path; the fused path ran; factor below 1 when clipping, 1 otherwise
TestZero3CpuAdamCopyBack (2 ranks; ratio 1.0 and 0.5)	staging pieces of 1,000 elements reuse both buffers many times per step; parameters bitwise equal to the separate passes; the fused copy-back ran

Hardware: 1 node of 8 H200 (143,771 MiB, driver 580.126.09), torch 2.8.0+cu126, pytest 9.1.1:

pytest unit/ops/adam/test_cpu_adam.py                                  # 137 passed, 5 skipped (master: 97 passed, 5 skipped)
pytest unit/v1/zero/test_zero.py::TestZero3CpuAdamGradScale            # 4 passed
pytest unit/v1/zero/test_zero.py::TestZero3CpuAdamCopyBack             # 2 passed
pytest unit/v1/zero/test_zero.py::TestZero3OffloadPinnedCopyBack       # 2 passed
pytest unit/runtime/zero/test_zero_offloadpp.py                        # 3 passed (master: 3 passed)

The 5 skips are the file's own fp16 cases. pre-commit run --files passes on every changed file.

Compatibility

Cell	Validated
ZeRO-3, CPU optimizer offload, pin_memory on, bf16, all of the optimizer on the CPU (ratio 1.0) and half (0.5)	yes (tests; 1 and 8 nodes)
AVX512 build (Intel Xeon Platinum 8488C, the H200 nodes)	yes
NEON / SVE builds, AVX2-only CPUs	no (AVX2 has no bf16 store, so copy_supported() is false and the copy keeps the old path; the clip factor is fused on every build)
fp16, GradScaler, NVMe offload, ZenFlow, Muon, SuperOffload, pin_memory off	path unchanged by design
ZeRO-1 / ZeRO-2	not touched

AI assistance

This change, its tests and this description were prepared with AI assistance (Claude).