DeepSpeed
7dfa67ae - fix(zero3): async grad offload + pinned offload buffers by default (#8207)

Commit
7 days ago
fix(zero3): async grad offload + pinned offload buffers by default (#8207) ## Problem ZeRO-3's gradient GPU→CPU offload in `partition_grads()` used a blocking `copy_()` without `non_blocking`, and its destination buffer defaulted to pageable host memory. This forced the offload onto a synchronous, low-bandwidth (staged pageable) path with no overlap against backward compute, even though the copy already runs on the dedicated `reduce_and_partition_stream`. ZeRO stage 1/2 already issues this copy with `non_blocking=True` (`stage_1_and_2.py:1530`); stage 3 is the inconsistent one. ## Changes - `offload_config.py`: default `offload_optimizer/offload_param.pin_memory` to `True`. Pinned (page-locked) host memory is required for async, full-bandwidth DMA; the prior `False` silently selected the slow staged pageable copy. Disable only on hosts with tight `ulimit -l` memlock. - `stage3.py`: issue the grad offload copy with `non_blocking=True`. Stays on `reduce_and_partition_stream`. - `stage3.py` (2nd commit): remove an orphaned helper (`async_inplace_copy_grad_to_fp32_buffer_from_gpu`) that referenced an uninitialized attribute and had no callers. The live stage 1/2 version is untouched. ## Validation 4× RTX 4080-SUPER, autotp=2, `offload_optimizer`, `cpu_adam`, per-rank CPU affinity: | model | baseline BWD | fixed BWD | Δ BWD | |----------------|--------------|-----------|-------| | Qwen2.5-1.5B | 1810 ms | 1308 ms | -28% | | Qwen2.5-3B | 3097 ms | 2540 ms | -18% | Memory footprint unchanged; FWD/STEP unchanged. --------- Signed-off-by: Guokai Ma <guokai.ma@intel.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Author
Parents
Loading