[hop free symbols][refactor] make create_graph_input always take example_value #138428

ydwu4 · 2024-10-20T18:21:26Z

Stack from ghstack (oldest at bottom):

[hop free symbols] lift free symbols in example_value when create_graph_input #138363
[hop free symbols] refactor tensor.to_list implementation to call wrap_fx_proxy. #139663
[hop free symbols][refactor] lift freevar to parent graph before lifting to subgraph #138559
[hop free symbols] replace ctx.save_for_backward to support symints/ints #138737
[hop free symbols][refactor] make map's save_for_backward to handle int #138558
-> [hop free symbols][refactor] make create_graph_input always take example_value #138428
[hop free symbols][refactor] make bound_symbols a dictionary #138345

Code refactoring only. We move the wrap_to_fake_tensor_logic out of wrap_fx_proxy for placeholders to provide the invariant that all graph inputs must set their example values when creating the inputs. This invariant helps us to identify all the free symbols in the graph in top-level and sub-graphs.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames @rec

…ple_value [ghstack-poisoned]

pytorch-bot · 2024-10-20T18:21:29Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/138428

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (3 Unrelated Failures)

As of commit 0a8e9df with merge base 8e27833 ():

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor_distributed, 1, 1, linux.g5.12xlarge.nvidia.gpu) (gh) (trunk failure)
distributed/_composable/test_replicate_with_compiler.py::ReplicateTest::test_bucketing_coalesced_op
inductor / cuda12.1-py3.10-gcc9-sm86 / test (inductor, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (trunk failure)
distributed/_composable/test_replicate_with_compiler.py::ReplicateTest::test_bucketing_coalesced_op
inductor / cuda12.1-py3.12-gcc9-sm86 / test (inductor, 1, 2, linux.g5.4xlarge.nvidia.gpu) (gh) (trunk failure)
distributed/_composable/test_replicate_with_compiler.py::ReplicateTest::test_bucketing_coalesced_op

This comment was automatically generated by Dr. CI and updates every 15 minutes.

…s take example_value" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames rec [ghstack-poisoned]

ydwu4 · 2024-10-20T18:46:54Z

torch/_dynamo/variables/builder.py

+        example_value = wrap_to_fake_tensor_and_record(
+            tensor_value,
+            tx=self.tx,
+            is_tensor=False,


is_tensor is dermined by whether target_cls in (TensorVariable, TensorWithTFOverrideVariable). NumpyNdarraryVaraible is not in them so is_tensor = False. This PR keeps this behavior unchanged and only do code-refactoring.

ydwu4 · 2024-10-20T18:47:07Z

torch/_dynamo/variables/builder.py

        # TODO: Maybe the tensor-ification should be built into the source,
        # rather than by special pattern match
+        example_value = wrap_to_fake_tensor_and_record(
+            wrapped_value, tx=self.tx, is_tensor=False, source=self.get_source()


same here for is_tensor=False

ydwu4 · 2024-10-20T18:47:18Z

torch/_dynamo/variables/builder.py

        options.update({"raw_value": value})

+        example_value = wrap_to_fake_tensor_and_record(
+            wrapped_value, tx=self.tx, is_tensor=False, source=self.get_source()


ydwu4 · 2024-10-20T19:01:29Z

torch/_dynamo/variables/builder.py

+        # the subgraph.
+        # See NOTE [HigherOrderOperator tracing design] for more details.
+
+        example_value = wrap_to_fake_tensor_and_record(


We move the wrap_to_fake_tensor_and_record logic in wrap_fx_proxy out, so that we could provide an example_value when create_graph_input. Since example_value has already been fakified, it won't get fakified again in wrap_fx_proxy.

…s take example_value" Code refactoring only. We move the wrap_to_fake_tensor_logic out of wrap_fx_proxy for placeholders to provide the invariant that **all graph inputs must have example values** cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames rec [ghstack-poisoned]

ydwu4 · 2024-10-21T21:56:54Z

test/dynamo/test_autograd_function.py


    class fwd_body_0(torch.nn.Module):
-        def forward(self, ctx, x: "f32[]", z: "f32[]", l_weird_b: "f32[]", l_weird_c: "f32[]"):
+        def forward(self, ctx : torch.autograd.function.Function, x: "f32[]", z: "f32[]", l_weird_b: "f32[]", l_weird_c: "f32[]"):


ctx's type is torch.autograd.funciton.Function instead of FunctionCtx because of https://github.com/pytorch/pytorch/blob/main/torch/_dynamo/side_effects.py#L270

What causes some of these to state their types but some of them not to?

…s take example_value" Code refactoring only. We move the wrap_to_fake_tensor_logic out of wrap_fx_proxy for placeholders to provide the invariant that **all graph inputs must set their example values when creating the inputs**. This invariant helps us to identify all the free symbols in the graph in top-level and sub-graphs. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames rec [ghstack-poisoned]

zou3519 · 2024-10-28T16:04:21Z

test/functorch/test_control_flow.py

                gm.cond_fn_0.code.strip(),
                """\
-def forward(self, l_iter_, l_x_, l_self_buffers_dec__cond_fn, l_self_modules_linear_parameters_bias__body_fn, l_self_modules_linear_parameters_weight__body_fn):
+def forward(self, l_iter_ : torch.Tensor, l_x_ : torch.Tensor, l_self_buffers_dec__cond_fn, l_self_modules_linear_parameters_bias__body_fn, l_self_modules_linear_parameters_weight__body_fn):


Why do the l_self_buffers_dec__cond_fn don't have types?

zou3519 · 2024-10-28T16:05:22Z

torch/_dynamo/output_graph.py

        if self.backward_state_proxy is None:
            if self.export:
                unimplemented("backward_state does not support export")
+            example_value = BackwardState()


Can you check (add a test?) that this doesn't cause a warning? Direct initialization of FunctionCtx usually induces a warniing

zou3519 · 2024-10-28T16:06:02Z

torch/_dynamo/output_graph.py

            proxy = self.root_tracer.create_graph_input(
                str(s0),
-                torch.SymInt,
+                type(s),


Is there a reason to pass type(s) if there's always an example value? create_graph_input can always get the type from the example value.

zou3519

Can you test that creating the backward state does not add user-facing warnings? We've had that problem in the past.

Other than that, this looks fine to me

…s take example_value" Code refactoring only. We move the wrap_to_fake_tensor_logic out of wrap_fx_proxy for placeholders to provide the invariant that **all graph inputs must set their example values when creating the inputs**. This invariant helps us to identify all the free symbols in the graph in top-level and sub-graphs. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames rec [ghstack-poisoned]

…nt (#138558) Pull Request resolved: #138558 Approved by: https://github.com/zou3519 ghstack dependencies: #138345, #138428

…nts (#138737) Pull Request resolved: #138737 Approved by: https://github.com/drisspg, https://github.com/zou3519, https://github.com/Chillee ghstack dependencies: #138345, #138428, #138558

…ing to subgraph (pytorch#138559) This refactoring is for getting a deterministic ordering of binding tensors and sizes of tensors. When seeing a free tensor x with shape (s0,) in subgraph, the ordering of lifting changes from ``` lift_x_in_child, lift_s0_in_child, lift_s0_in_parent, lift_x_in_parent ``` to ``` lift_x_in_parent, lift_s0_in_parent, lift_x_in_child, lift_s0_in_child ``` This produces a determinstic ordering of handling the symints in lifted tensors. This is also the current contract of dynamo top-level graph: we lift free_symbols in sizes after tensor x and insert the free symbols before the tensor x's proxy. Pull Request resolved: pytorch#138559 Approved by: https://github.com/zou3519 ghstack dependencies: pytorch#138345, pytorch#138428, pytorch#138558, pytorch#138737

…p_fx_proxy. (#139663) Refactoring only. Previously, we manually cal SymNodeVariable.create, now we handle it with wrap_fx_proxy. This unifies the handling of operations that produce symints in wrap_fx_proxy. Pull Request resolved: #139663 Approved by: https://github.com/zou3519 ghstack dependencies: #138345, #138428, #138558, #138737, #138559

…p_fx_proxy. (pytorch#139663) Refactoring only. Previously, we manually cal SymNodeVariable.create, now we handle it with wrap_fx_proxy. This unifies the handling of operations that produce symints in wrap_fx_proxy. Pull Request resolved: pytorch#139663 Approved by: https://github.com/zou3519 ghstack dependencies: pytorch#138345, pytorch#138428, pytorch#138558, pytorch#138737, pytorch#138559

[hop free symbols][refactor] make create_graph_input always take exam…

37a2f4c

…ple_value [ghstack-poisoned]

ydwu4 requested a review from zou3519 as a code owner October 20, 2024 18:21

ydwu4 mentioned this pull request Oct 18, 2024

[hop free symbols][refactor] make bound_symbols a dictionary #138345

Closed

ydwu4 mentioned this pull request Oct 19, 2024

[hop free symbols] lift free symbols in example_value when create_graph_input #138363

Closed

pytorch-bot bot added ciflow/inductor module: dynamo release notes: fx release notes category labels Oct 20, 2024

ydwu4 added 3 commits October 20, 2024 11:24

Update on "[hop free symbols][refactor] make create_graph_input alway…

058ba20

…s take example_value" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames rec [ghstack-poisoned]

Update on "[hop free symbols][refactor] make create_graph_input alway…

8685741

…s take example_value" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames rec [ghstack-poisoned]

Update on "[hop free symbols][refactor] make create_graph_input alway…

28e9ea2

…s take example_value" cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames rec [ghstack-poisoned]

ydwu4 commented Oct 20, 2024

View reviewed changes

ydwu4 added topic: not user facing topic category and removed release notes: fx release notes category labels Oct 20, 2024

ydwu4 added 2 commits October 21, 2024 13:19

ydwu4 commented Oct 21, 2024

View reviewed changes

This was referenced Oct 22, 2024

[hop free symbols][refactor] make map's save_for_backward to handle int #138558

Closed

[hop free symbols][refactor] lift freevar to parent graph before lifting to subgraph #138559

Closed

ydwu4 requested review from anijain2305 and ezyang October 23, 2024 18:19

ydwu4 mentioned this pull request Oct 23, 2024

[hop free symbols] replace ctx.save_for_backward to support symints/ints #138737

Closed

ezyang approved these changes Oct 25, 2024

View reviewed changes

ydwu4 added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 25, 2024

zou3519 reviewed Oct 28, 2024

View reviewed changes

zou3519 approved these changes Oct 28, 2024

View reviewed changes

ydwu4 added 2 commits October 28, 2024 15:48

This was referenced Nov 4, 2024

[hop free symbols] refactor tensor.to_list implementation to call wrap_fx_proxy. #139663

Closed

tmp #139664

Closed

pytorchmergebot closed this in dc3a6a9 Nov 4, 2024

pytorchmergebot added the Merged label Nov 4, 2024

pytorchmergebot pushed a commit that referenced this pull request Nov 4, 2024

[hop free symbols][refactor] make map's save_for_backward to handle i…

ac20d0f

…nt (#138558) Pull Request resolved: #138558 Approved by: https://github.com/zou3519 ghstack dependencies: #138345, #138428

github-actions bot deleted the gh/ydwu4/162/head branch December 5, 2024 02:14

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[hop free symbols][refactor] make create_graph_input always take example_value #138428

[hop free symbols][refactor] make create_graph_input always take example_value #138428

Uh oh!

ydwu4 commented Oct 20, 2024 •

edited

Loading

Uh oh!

pytorch-bot bot commented Oct 20, 2024 •

edited

Loading

Uh oh!

ydwu4 Oct 20, 2024

Uh oh!

ydwu4 Oct 20, 2024

Uh oh!

ydwu4 Oct 20, 2024

Uh oh!

ydwu4 Oct 20, 2024 •

edited

Loading

Uh oh!

ydwu4 Oct 21, 2024 •

edited

Loading

Uh oh!

zou3519 Oct 28, 2024

Uh oh!

zou3519 Oct 28, 2024

Uh oh!

zou3519 Oct 28, 2024

Uh oh!

zou3519 Oct 28, 2024

Uh oh!

zou3519 left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

[hop free symbols][refactor] make create_graph_input always take example_value #138428

[hop free symbols][refactor] make create_graph_input always take example_value #138428

Uh oh!

Conversation

ydwu4 commented Oct 20, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Oct 20, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/138428

✅ You can merge normally! (3 Unrelated Failures)

Uh oh!

ydwu4 Oct 20, 2024

Choose a reason for hiding this comment

Uh oh!

ydwu4 Oct 20, 2024

Choose a reason for hiding this comment

Uh oh!

ydwu4 Oct 20, 2024

Choose a reason for hiding this comment

Uh oh!

ydwu4 Oct 20, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ydwu4 Oct 21, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zou3519 Oct 28, 2024

Choose a reason for hiding this comment

Uh oh!

zou3519 Oct 28, 2024

Choose a reason for hiding this comment

Uh oh!

zou3519 Oct 28, 2024

Choose a reason for hiding this comment

Uh oh!

zou3519 Oct 28, 2024

Choose a reason for hiding this comment

Uh oh!

zou3519 left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ydwu4 commented Oct 20, 2024 •

edited

Loading

pytorch-bot bot commented Oct 20, 2024 •

edited

Loading

ydwu4 Oct 20, 2024 •

edited

Loading

ydwu4 Oct 21, 2024 •

edited

Loading