DeepSpeed
d3dcf9ad - Extend unmanaged gradient accumulation to ZeRO offload

Commit
26 days ago
Extend unmanaged gradient accumulation to ZeRO offload Allow managed_gradient_accumulation=False with ZeRO optimizer-state and parameter offload (CPU/NVMe). Stage 2/3 defer boundary norm and FP32/NVMe gradient copy into finalize_gradient_accumulation_boundary(); stage 1 continues to finalize via allreduce_gradients() at step(). Docs and equivalence tests cover stages 1-3 optimizer offload and stage-3 param offload. Validated on a 2-GPU node (full -k Unmanaged suite, 31 passed) incl. optimizer offload [1,2,3] and param+optimizer offload stage 3. Signed-off-by: Olatunji Ruwase <tunji.ruwase@snowflake.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Parents
Loading