Fix DeepCompile ZeRO-3 gathered parameter ownership#8157
Draft
tohtana wants to merge 1 commit into
Draft
Conversation
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
DeepCompile's ZeRO-3 eager fallback could gather a partitioned parameter while Dynamo was only evaluating guards, carry an unowned full parameter into a later forward, or keep tracking a parameter after an explicit
GatheredParameterscontext had adopted it. Those paths gave the fallback and user code conflicting responsibility for the same gathered lifetime.Why it matters
Guard evaluation must not change parameter residency. At each forward boundary, a nonpersistent full parameter must be partitioned unless an explicit user context owns it, and exactly one owner must decide when it is safe to release.
Solution
This change suppresses eager fallback gathers during Dynamo guard evaluation, releases unowned available parameters before the next outermost forward, and records one fallback owner per gathered parameter. Entering
GatheredParameterstransfers ownership to the user context, including nested and exceptional exits, so the fallback cannot repartition user-owned storage. This branch is the stack base and does not change compiler-global lifecycle or scheduler policy.