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".

[pengdurice]

let's import and use _FALLBACK_OWNER_ATTR?

[tohtana]

Good catch! Fixed.

[pengdurice]

is "torch._dynamo.guards" prone to change e.g. when torch version changes?

[tohtana]

Yes, it can change. I could not find a public API for detecting guard evaluation specifically. torch.compiler.is_compiling() returns false in this path. However, it still works even we fail to find torch._dynamo.guards.

If the module path changes, this check returns false and the guard access is treated as a normal eager fallback access. That causes an unnecessary all-gather and may increase communication and memory usage, but it does not skip an all-gather required by actual eager execution.

[pengdurice]

General question, if some one does this:

loss = engine(input)
with GatheredParameters([W]):
    log(W.norm())
 engine.backward(loss)

will that cause issue since the GatheredParameters context manager already partitions the parameters when it exits?

[tohtana]

@pengdurice Great catch! Yes, in this pattern, the parameter can be partitioned too early.

I updated the PR branch to address it. Instead of transferring the ownership of the fallback and the gathered parameter, we now records claims of the ownership as a set (it is like a reference count). The parameter is partitioned only after the last claim is released.