[export] Update dynamo_graph_capture_for_export to return GraphModule. #166091

zhxchen17 · 2025-10-22T21:05:22Z

Make dynamo_graph_capture_for_export return a more compatible GraphModule object which is closer the the original behavior of dynamo

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @msaroufim @dcci @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames @Lucaskabela

pytorch-bot · 2025-10-22T21:05:26Z

🔗 Helpful Links

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

📄 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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

All Docker build failed due to Ubuntu archive outage

❌ 2 New Failures, 1 Unrelated Failure

As of commit 3e31eb6 with merge base 1e836bc ():

NEW FAILURES - The following jobs have failed:

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

trunk / linux-jammy-py3-clang12-executorch / test (executorch, 1, 1, lf.linux.2xlarge, unstable) (gh) (#166072)
examples/models/llama3_2_vision/text_decoder/test/test_text_decoder.py::TextDecoderTest::test_llama3_2_text_decoder_aoti

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

linux-foundation-easycla · 2025-10-22T21:05:28Z

The committers listed above are authorized under a signed CLA.

✅ login: zhxchen17 / name: Zhengxu Chen (3e31eb6)

tugsbayasgalan · 2025-10-23T17:49:55Z

torch/fx/graph.py

        else:
            return "\n    " + "".join(x + "; " for x in has_annotation) + "\n"

+    def gen_var_bindings(self, fn_args, free_vars, expanded_def) -> str:


Is it just codemod?

I added a subclass called _ExportCodegen which adds shuffling.

tugsbayasgalan · 2025-10-23T17:50:14Z

torch/export/_trace.py

-            gm_torch_level._in_spec,
-            out_spec,
-        )
+    gm_torch_level.graph._codegen.pytree_info = _PyTreeInfo(


Hmm why did it change?

We will have different subclass called _ExportCodeGen here, so it's wrong to always assign _PyTreeCodegen here.

test/export/test_experimental.py

test/distributed/tensor/test_dtensor_export.py

torch/_dynamo/functional_export.py

zhxchen17 · 2025-10-23T20:46:39Z

@anijain2305 I tested this with autoparallel and pytorch unittests.

Now with diff we can make dynamo -> aot autograd working without:

install_free_tensors
restore_state_dict

My plan is to codemod the API usage in autoparallel first. Then I will go for 6lib.

test/export/test_experimental.py

tugsbayasgalan

Looks mostly good. Pls get review from @anijain2305 as well.

test/export/test_experimental.py

tugsbayasgalan · 2025-10-24T05:43:44Z

test/export/test_experimental.py

+        def forward(self, args_0, args_1):
+            _tree_leaf_0, _tree_leaf_1, _tree_leaf_2, = pytree.tree_leaves((self, args_0, args_1,))
+            L_fw_in_ , L_bw_in_ , = self._in_shuffle_graph(_tree_leaf_0, _tree_leaf_1, _tree_leaf_2)
+            l_fw_in_ = L_fw_in_


nit: this feels little redundant, any chance this can be simplified?

what do you mean specifically by redundant?

test/export/test_experimental.py

tugsbayasgalan · 2025-10-27T17:31:12Z

torch/_dynamo/functional_export.py

-        return types.MethodType(pytree_call, mod.__self__)
-    else:
-        return pytree_call
+def normalize_graph_module(gm):


I am pretty sure this is not enough for export, but probably fine for now.

yeah, let's tackle torch.export separately.

tugsbayasgalan · 2025-10-27T17:33:40Z

torch/_dynamo/functional_export.py

+        ]
+        graph_module.graph._codegen = _ExportCodeGen(
+            _PyTreeInfo(
+                argument_names(inspect.signature(mod), args, kwargs),


Do we actually use argument names in torch IR produced graph? If not, can we add this logic in the follow up PR? if the previous export behaviour is that it doesn't produce torch IR with correct user argument names, i am ok btw.

ok I can add a TODO here if you're ok with that.

tugsbayasgalan · 2025-10-27T17:34:16Z

torch/_dynamo/functional_export.py

+        assert not hasattr(graph_module, "_out_shuffle_graph")
+        graph_module._in_shuffle_graph = pyt.in_shuffle_graph
+        graph_module._out_shuffle_graph = pyt.out_shuffle_graph
+        delattr(graph_module, "_param_name_to_source")


Hmm is this the issue in previous torch IR graph capture as well?

I think in previous impl, you just transform and return a new graph module which doesn't have all the attached attributes like this.

tugsbayasgalan · 2025-10-27T17:40:13Z

torch/_dynamo/functional_export.py

+        graph_module._out_shuffle_graph = pyt.out_shuffle_graph
+        delattr(graph_module, "_param_name_to_source")
+        graph_module.recompile()
+        graph_module.meta["module_call_specs"] = (


Just for my understanding, technically you don't need this right because you are anyway running the bytecode?

technically I'm not using this. This is just for passing some basic torch.export tests.

zhxchen17 · 2025-10-27T22:39:35Z

@pytorchbot merge

pytorchmergebot · 2025-10-27T22:41:30Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2025-10-27T22:57:35Z

Merge failed

Reason: 1 jobs have failed, first few of them are: trunk / inductor-build / build

Details for Dev Infra team

Raised by workflow job

zhxchen17 · 2025-10-27T23:07:51Z

@pytorchbot merge -i

pytorchmergebot · 2025-10-27T23:09:53Z

Merge started

Your change will be merged while ignoring the following 1 checks: trunk / inductor-build / build

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2025-10-27T23:15:29Z

Merge failed

Reason: 1 jobs have failed, first few of them are: trunk / linux-jammy-cuda12.8-py3.10-gcc11 / test (distributed, 1, 3, lf.linux.g4dn.12xlarge.nvidia.gpu)

Details for Dev Infra team

Raised by workflow job

zhxchen17 · 2025-10-28T04:00:05Z

@pytorchbot merge -i

Test failures seem not related.

pytorchmergebot · 2025-10-28T04:01:52Z

Merge started

Your change will be merged while ignoring the following 3 checks: trunk / linux-jammy-py3-clang12-executorch / test (executorch, 1, 1, lf.linux.2xlarge, unstable), trunk / linux-jammy-cuda12.8-py3.10-gcc11 / test (default, 3, 5, lf.linux.g6.4xlarge.experimental.nvidia.gpu), trunk / linux-jammy-cuda12.8-py3.10-gcc11 / test (default, 1, 5, lf.linux.g6.4xlarge.experimental.nvidia.gpu)

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

zhxchen17 requested review from angelayi, avikchaudhuri, tugsbayasgalan and ydwu4 as code owners October 22, 2025 21:05

pytorch-bot bot added ciflow/inductor module: dynamo oncall: distributed Add this issue/PR to distributed oncall triage queue labels Oct 22, 2025

pytorch-bot bot added the release notes: export label Oct 22, 2025

facebook-github-bot added the fx label Oct 22, 2025

zhxchen17 added topic: not user facing topic category and removed fx release notes: export labels Oct 22, 2025

zhxchen17 force-pushed the zhxchen17/precompile/export_gm branch 4 times, most recently from c5e4a57 to 4331a46 Compare October 23, 2025 17:14

tugsbayasgalan reviewed Oct 23, 2025

View reviewed changes

zhxchen17 force-pushed the zhxchen17/precompile/export_gm branch 2 times, most recently from 99ea1aa to 1b303ae Compare October 23, 2025 19:36

zhxchen17 requested a review from anijain2305 October 23, 2025 20:09

zhxchen17 force-pushed the zhxchen17/precompile/export_gm branch 2 times, most recently from a7bf577 to 031371e Compare October 23, 2025 20:42

zhxchen17 requested a review from tugsbayasgalan October 23, 2025 20:43

zhxchen17 force-pushed the zhxchen17/precompile/export_gm branch from 031371e to 8b81d5b Compare October 24, 2025 03:10

tugsbayasgalan reviewed Oct 24, 2025

View reviewed changes

test/export/test_experimental.py Show resolved Hide resolved

tugsbayasgalan reviewed Oct 24, 2025

View reviewed changes

anijain2305 reviewed Oct 24, 2025

View reviewed changes

test/export/test_experimental.py Show resolved Hide resolved

zhxchen17 force-pushed the zhxchen17/precompile/export_gm branch 2 times, most recently from a0fd1b6 to b43dcec Compare October 24, 2025 18:33

zhxchen17 requested a review from anijain2305 October 24, 2025 18:43

tugsbayasgalan reviewed Oct 27, 2025

View reviewed changes

zhxchen17 force-pushed the zhxchen17/precompile/export_gm branch from b43dcec to b251c24 Compare October 27, 2025 19:15

tugsbayasgalan approved these changes Oct 27, 2025

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 27, 2025

pytorchmergebot added the merging label Oct 27, 2025

pytorchmergebot removed the merging label Oct 27, 2025

pytorchmergebot added the merging label Oct 27, 2025

pytorchmergebot removed the merging label Oct 27, 2025

zhxchen17 force-pushed the zhxchen17/precompile/export_gm branch from b251c24 to 6f5b716 Compare October 28, 2025 01:26

[export] Update dynamo_graph_capture_for_export to return GraphModule.

3e31eb6

zhxchen17 force-pushed the zhxchen17/precompile/export_gm branch from 6f5b716 to 3e31eb6 Compare October 28, 2025 01:43

pytorchmergebot added the merging label Oct 28, 2025

pytorchmergebot closed this in f93ea7d Oct 28, 2025

pytorchmergebot added Merged and removed merging labels Oct 28, 2025

Uh oh!

[export] Update dynamo_graph_capture_for_export to return GraphModule. #166091

[export] Update dynamo_graph_capture_for_export to return GraphModule. #166091

Conversation

zhxchen17 commented Oct 22, 2025 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Oct 22, 2025 • 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/166091

❗ 1 Active SEVs

❌ 2 New Failures, 1 Unrelated Failure

Uh oh!

linux-foundation-easycla bot commented Oct 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zhxchen17 commented Oct 23, 2025

Uh oh!

Uh oh!

tugsbayasgalan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhxchen17 commented Oct 27, 2025

Uh oh!

pytorchmergebot commented Oct 27, 2025

Merge started

Uh oh!

pytorchmergebot commented Oct 27, 2025

Merge failed

Uh oh!

zhxchen17 commented Oct 27, 2025

Uh oh!

pytorchmergebot commented Oct 27, 2025

Merge started

Uh oh!

pytorchmergebot commented Oct 27, 2025

Merge failed

Uh oh!

zhxchen17 commented Oct 28, 2025

Uh oh!

pytorchmergebot commented Oct 28, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

zhxchen17 commented Oct 22, 2025 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Oct 22, 2025 •

edited

Loading

linux-foundation-easycla bot commented Oct 22, 2025 •

edited

Loading