Skip to content

Make DeepCompile ZeRO-3 scheduling memory-aware#8129

Draft
tohtana wants to merge 5 commits into
deepspeedai:masterfrom
tohtana:tohtana/fix/deepcompile-zero3-memory-aware-scheduler-review-fixes
Draft

Make DeepCompile ZeRO-3 scheduling memory-aware#8129
tohtana wants to merge 5 commits into
deepspeedai:masterfrom
tohtana:tohtana/fix/deepcompile-zero3-memory-aware-scheduler-review-fixes

Conversation

@tohtana

@tohtana tohtana commented Jul 10, 2026

Copy link
Copy Markdown
Collaborator

This PR improves DeepCompile's scheduling and robustness.

  • Make DeepCompile ZeRO-3 scheduling memory-aware by deriving a data-parallel-group-consistent gather budget from conservative per-node memory profiles and using it to order parameter gathers and releases.
  • Preserve ZeRO parameter ownership and compiled-backward lifecycle state across eager fallback, compilation failures, and overlapping DeepSpeed engines.
  • Add rank-consistency diagnostics, focused regression coverage, and documentation explaining the scheduling model and its safety constraints.

tohtana added 3 commits July 16, 2026 13:39
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
@tohtana
tohtana force-pushed the tohtana/fix/deepcompile-zero3-memory-aware-scheduler-review-fixes branch from 7e011b0 to d219da5 Compare July 16, 2026 21:46
tohtana added 2 commits July 16, 2026 15:51
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant