github-featured / deepspeedai--deepspeed

发布

0 发布 128 Git 标签

RSS 订阅源
  • v0.19.5 cf44300453
    比较

    Unmanaged gradient accumulation: ZeRO offload support (#8225)

    frostbyte_neo 发布于 2026-08-09 18:28:04 +00:00

    Summary

    • Extends unmanaged gradient accumulation
      (managed_gradient_accumulation=false) to ZeRO optimizer-state and
      parameter offload       (CPU/NVMe). Follow-up to #8217 (ZeRO stage 3, now
      merged).
    • Stage 2/3: grads still reduce/partition every backward(); step()
      finalizes deferred offload boundary work (grad norms + FP32/NVMe copy)
      via finalize_gradient_accumulation_boundary().
    • Stage 1: continues to reduce at step() via allreduce_gradients(),
      which already performs offload boundary finalization when the boundary
      flag is true.
    • Pipeline parallelism, DeepCompile, Apex AMP, and stage-0/1
      overlap_comm remain unsupported.

    Test plan

    Validated on a 2-GPU node:

    • Full -k Unmanaged suite (31 passed), including:
      • test_unmanaged_matches_managed_optimizer_offload[1|2|3]
      • test_unmanaged_matches_managed_param_offload (stage 3)
    • existing non-offload unmanaged equivalence / varying-GAS / rejection
      tests
    • Docs updated (config-json.md, training.rst); previewable on
      rtd-staging

    Made with Cursor

    Signed-off-by: Olatunji Ruwase tunji.ruwase@snowflake.com
    Co-authored-by: Cursor cursoragent@cursor.com
    Co-authored-by: Ma, Guokai guokai.ma@gmail.com

    下载附件