Skip to content

DeepCompile: stabilize ZeRO-3 parameter guards - #8328

Draft
tohtana wants to merge 1 commit into
deepspeedai:masterfrom
tohtana:tohtana/deepcompile-zero3-parameter-guards
Draft

DeepCompile: stabilize ZeRO-3 parameter guards#8328
tohtana wants to merge 1 commit into
deepspeedai:masterfrom
tohtana:tohtana/deepcompile-zero3-parameter-guards

Conversation

@tohtana

@tohtana tohtana commented Aug 27, 2026

Copy link
Copy Markdown
Collaborator

Under DeepCompile's ZeRO-3, the same module parameter can be empty (while partitioned) and full-sized param.data (while gathered). Dynamo could treat the state transition as a change to the model and recompile the graph.

The eager fallback added in #8059 gathers ZeRO-3 parameters for frames that Dynamo skips, which made these temporary state changes visible during tracing. #8157 prevented one parameter lookup performed during guard evaluation from triggering a gather. However, Dynamo could still build tensor guards from the parameter’s temporary gathered or partitioned representation, and ZeROOrderedDict could still read the current ds_status while Dynamo was compiling.

This PR makes Dynamo's graph reuse check compare the parameter's identity and logical properties, including its type, dispatch keys, dtype, device, layout, requires_grad, ds_id, and ds_shape. It does not make the check depend on the temporary param.data shape or ds_status. As the result, DeepCompile can keep Dynamo's existing tensor checks unchanged for non-ZeRO tensors.
When DeepCompile supplies a full-sized placeholder during tracing, it now looks up the corresponding module parameter before building the check. It also returns parameters from ZeROOrderedDict during compilation without reading ds_status.

Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>
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.

1 participant