Skip to content

Fix DeepCompile profiling memory cleanup - #8106

Merged
PKUWZP merged 2 commits into
deepspeedai:masterfrom
tohtana:tohtana/fix/deepcompile-zero3-oom-memory-upstream
Jul 1, 2026
Merged

Fix DeepCompile profiling memory cleanup#8106
PKUWZP merged 2 commits into
deepspeedai:masterfrom
tohtana:tohtana/fix/deepcompile-zero3-oom-memory-upstream

Conversation

@tohtana

@tohtana tohtana commented Jul 1, 2026

Copy link
Copy Markdown
Collaborator

This PR reduces DeepCompile ZeRO-3 memory pressure in profiling and compiled backward by fixing three related issues:

  • Let the partitioner keep using its min-cut rematerialization policy for non-parameter activations, while still forcing ZeRO-3 parameter aliases/casts to be recomputed.
  • Compute selective-gather persistence capacity from profiled transient headroom instead of current allocator availability.
  • Avoid retaining temporary profiling outputs across warmup/measured profiling iterations, and clear gathered ZeRO-3 parameters when profiling exits through an exception.
  • Mark incomplete memory profiles explicitly, synchronize the completion status across ranks, and skip profile-dependent prefetch/selective-gather decisions when profiling data is incomplete.

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
@tohtana
tohtana requested review from loadams and tjruwase as code owners July 1, 2026 18:43
@tohtana
tohtana requested review from PKUWZP, eternalNight, sfc-gh-truwase and tjruwase and removed request for loadams and tjruwase July 1, 2026 18:43

@PKUWZP PKUWZP left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, approved.

@PKUWZP
PKUWZP added this pull request to the merge queue Jul 1, 2026
Merged via the queue into deepspeedai:master with commit 131779c Jul 1, 2026
13 of 17 checks passed
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.

2 participants