feat(zenflow): run the overlapped CPU optimizer in a native process (#8058)
## What changes this PR introduce
ZenFlow's overlapped CPU optimizer step previously ran in a Python
`multiprocessing`
subprocess coordinated by a pickling `Pipe`. This PR moves that
optimizer into a
**native CPU optimizer process** packaged inside the `cpu_adam` op,
coordinated through a
**shared-memory POSIX-semaphore control block** instead of pickling. The
Adam state is
allocated in that process, NUMA-local to the optimizer's pinned thread
pool.
Highlights:
- **Fused multi-tensor CPU Adam** (`adam_update_multi`): drives a whole
flattened partition
in C++ and writes the stale snapshot natively, removing the
per-parameter Python↔C++ loop
and the Python-side `clone()`.
- **`ZenFlowAdam`** native class: a pinned `std::thread` pool (pinned to
ZenFlow's dedicated
cores) running the serial Adam kernel per slice, driven from the main
process via the
shared-memory control block (`run_worker` / `submit` / `wait`).
- Covers ZeRO **stages 1, 2, and 3**; removes the old pickling
subprocess entirely.
- **Chunked copyback**: streaming the updated fp32 master partition back
to the GPU bit16
partition in chunks drops a transient GPU spike from ~2944 MiB to ~256
MiB for a
0.75B-param partition (the old `fp32.to(device)` materialized the whole
fp32 partition on
the GPU first).
- `ZenFlowCPUAdam` is now a **recognized ZeRO optimizer**, so
`zero_allow_untested_optimizer`
is no longer required in ZenFlow configs.
## Correctness & performance
- Bit-identical to the previous subprocess path: cross-process and
fused-op unit tests, plus
seeded end-to-end loss across ZeRO stages 1/2/3.
- Real Qwen2.5-1.5B end-to-end (ZeRO-2, CPU offload, 2 GPUs): per-step
throughput unchanged
(no regression). Small / IPC-bound configurations are faster (the
per-step pickling/IPC
overhead is removed).
## Dependency / merge order
This branch is based on top of #7771 ("Fix ZenFlow NaN under
PyTorch-style backward"), so its
`backward_prologue` commit rides along here. **Please merge #7771 first,
then this PR** — after
#7771 lands, that commit is already in `master` and only the
native-optimizer changes remain.
## Testing
- `tests/unit/ops/adam/test_cpu_adam.py`:
`test_zenflow_adam_cross_process` (production path,
bit-identical to the fused reference) and `TestCPUAdamFusedMultiTensor`.
- End-to-end ZenFlow finetuning on ZeRO stages 1/2/3 (single- and
multi-GPU).
Note: the native optimizer process uses POSIX semaphores and is
Linux-only.
---------
Signed-off-by: Tingfeng Lan <erc8gx@virginia.edu>