DeepSpeed
1ddbe484 - feat(zenflow): run the overlapped CPU optimizer in a native process (#8058)

Commit
5 days ago
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>
Author
Parents
Loading