Fix DeepCompile ZeRO-3 gathered parameter ownership (#8157)
## 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.
---------
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
