## 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>
Ma, Guokai
7dfa67aeaf