Fix DeepCompile ZeRO-3 gathered parameter ownership

[tohtana]

Problem

DeepCompile inserts ZeRO-3 parameter all-gather and release operations into compiled graphs. When Dynamo skips a frame because of a graph break, however, that frame executes eagerly and does not run those graph operations. The eager fallback introduced in #8059 handles this case by all-gathering a partitioned parameter when the skipped frame accesses it through ZeROOrderedDict.

The fallback is enabled around DeepSpeedEngine.forward(). Dynamo guard evaluation occurs inside that outer forward context and also resolves parameters through ZeROOrderedDict, while torch.compiler.is_compiling() is false. The fallback could therefore mistake a guard lookup for actual eager execution and unnecessarily all-gather the parameter.

Parameters gathered by the fallback are normally partitioned after backward, but that cleanup does not run when backward is skipped. A fallback-gathered parameter may also be passed to an explicit GatheredParameters context, which must keep the full tensor available until the context exits.

Why it matters

These cases require different behavior:

• Dynamo guard evaluation should not trigger an all-gather.
• A parameter gathered for an eagerly executed frame must remain available through backward and then be partitioned.
• If backward does not run, a leftover full parameter must be partitioned before the next outermost forward.
• A parameter covered by GatheredParameters must remain fully gathered until that context exits.

Without distinguishing these cases, a full parameter can remain allocated into a later forward, or fallback cleanup can partition it while a GatheredParameters block is still using it.

Solution

This PR:

• detects parameter access during Dynamo guard evaluation and skips the eager fallback all-gather;
• partitions leftover nonpersistent full parameters before the next outermost forward when the normal post-backward cleanup did not run;
• removes a parameter from fallback cleanup when it is passed to GatheredParameters, so that context alone partitions it on exit;
• restores the GatheredParameters state even when context exit raises; and
• rejects nested GatheredParameters contexts that overlap on the same parameter, while continuing to allow nesting over disjoint parameter sets.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 4563ff10f8

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".