change torch._dynamo.export(aten_graph=...) to allow pre_autograd tracing #98031

bdhirsh · 2023-03-30T22:11:41Z

pre_autograd tracing is still early, but it should work for basic cases. This PR changes the API a bit for export to expose pre_autograd tracing. Name bikeshedding is welcome, but it looks like:

torch._dynamo.export(..., aten_graph="aten_pre_autograd")

cc @soumith @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire @andrewor14

Stack from ghstack (oldest at bottom):

cc @soumith @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire

…cing [ghstack-poisoned]

pytorch-bot · 2023-03-30T22:11:44Z

🔗 Helpful Links

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

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

✅ No Failures

As of commit 7f6bdb5:
💚 Looks good so far! There are no failures yet. 💚

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

…cing ghstack-source-id: 7c13ae9d8f47d4bb85c7ab5b7b49c3cad7ec6fd9 Pull Request resolved: #98031

guangy10 · 2023-03-31T16:59:07Z

torch/_dynamo/eval_frame.py

-        aten_graph (bool): If True, exports a graph with ATen operators.
-        If False, exports a graph with Python operators. Default is False.
+        aten_graph (str): Valid options include:
+          "none": export a graph with Python operations. Default is False.


Suggested change

"none": export a graph with Python operations. Default is False.

"none": export a graph with Python operations.

guangy10 · 2023-03-31T17:00:10Z

torch/_dynamo/eval_frame.py

+          "aten": export a graph with ATen operations.
+          "aten_pre_autograd": export a graph with pre_autograd ATen operations.


What's the difference between pre_autograd aten ops and non-pre_autograd ones?

ezyang · 2023-04-03T14:50:44Z

torch/_dynamo/eval_frame.py

@@ -584,7 +584,7 @@ class directly; instead, use :func:`torch._export.dynamic_dim`.
 def export(
    f: Callable[..., Any],
    *args,
-    aten_graph: bool = False,
+    aten_graph: str = "none",


I don't think I would have personally gone ahead and done a BC-breaking change here. But if we are in the business of making a BC-breaking change, I think we ought to also rename the keyword argument as well. It's weird and redundant to say aten_graph="aten". How about dialect or target_dialect? Also cc @SherlockNoMad for opinions.

dialect sounds good to me (but interested in Sherlock's opinion).

I'm also concerned about BC risk - if you're worried about changing the bool, I'm happy to leave it alone (although two bools to control the "dialect" feels bad). Fixing any internal call-sites will be annoying but doable.

If you're willing to put in the effort, fixing the API earlier rather than later is good. I just remember I had specifically asked you to plumb this, and had been thinking you'd do the low effort thing.

you're right - in interesting of landing sooner, I'm avoiding the API concern for now and adding an extra boolean argument. cc @andrewor14

…utograd tracing" pre_autograd tracing is still early, but it should work for basic cases. This PR changes the API a bit for export to expose pre_autograd tracing. Name bikeshedding is welcome, but it looks like: ``` torch._dynamo.export(..., aten_graph="aten_pre_autograd") ``` cc andrewor14 cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

bdhirsh · 2023-04-11T15:13:02Z

updated, fixed tests

ezyang · 2023-04-11T19:07:17Z

torch/fx/experimental/proxy_tensor.py

@@ -414,7 +414,8 @@ def can_handle_tensor(x):
    else:
        constant = None

-    track_tensor_tree(out, proxy_out, constant=constant, tracer=tracer)
+    with inside_mode(proxy_mode):


This seems surprising to me.

still waiting on explainer

sigh, wrote the comment when I initially made this update and forgot to submit it.

I'd be happy to make this a separate PR. The issue was that:

(1) We detach tensors as part of creating proxies.

(2) In pre_autograd tracing, we currently push TorchProxyDispatchMode onto both the autograd mode stack, and the original python key mode stack

(3) The detach() calls get intercepted by the proxy mode on the original python key mode stack.

So we need to be careful that any aten ops that we call inside of TorchProxyDispatchMode happen in a with inside_mode() now.

I think we agreed 2-3 weeks ago that we could avoid re-entrant issues if we solved the "fallthrough keys can't be intercepted by the python dispatcher" issue, but we agreed this is difficult.

Ok put this comment in the code?

ezyang · 2023-04-11T19:08:03Z

This looks fine EXCEPT for the one line at the end, how come?

…utograd tracing" pre_autograd tracing is still early, but it should work for basic cases. This PR changes the API a bit for export to expose pre_autograd tracing. Name bikeshedding is welcome, but it looks like: ``` torch._dynamo.export(..., aten_graph="aten_pre_autograd") ``` cc andrewor14 cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

…utograd tracing" pre_autograd tracing is still early, but it should work for basic cases. This PR changes the API a bit for export to expose pre_autograd tracing. Name bikeshedding is welcome, but it looks like: ``` torch._dynamo.export(..., aten_graph="aten_pre_autograd") ``` cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire andrewor14 cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

bdhirsh · 2023-04-24T19:40:07Z

@pytorchbot merge

pytorchmergebot · 2023-04-24T19:42:38Z

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 · 2023-04-24T19:47:43Z

Merge failed

Reason: 1 jobs have failed, first few of them are: trunk / macos-12-py3-arm64 / build

Details for Dev Infra team

Raised by workflow job

…utograd tracing" pre_autograd tracing is still early, but it should work for basic cases. This PR changes the API a bit for export to expose pre_autograd tracing. Name bikeshedding is welcome, but it looks like: ``` torch._dynamo.export(..., aten_graph="aten_pre_autograd") ``` cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire andrewor14 cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

bdhirsh · 2023-04-25T19:19:13Z

@pytorchbot merge

pytorchmergebot · 2023-04-25T19:21:13Z

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

change torch._dynamo.export(aten_graph=...) to allow pre_autograd tra…

d7e356b

…cing [ghstack-poisoned]

bdhirsh mentioned this pull request Mar 30, 2023

aot_autograd: avoid using intermediate_base logic unnecessarily #97786

Closed

bdhirsh mentioned this pull request Mar 30, 2023

fix per-dispatchkey-mode caching bug #98030

Closed

github-actions bot requested review from albanD, antoniojkim, ezyang, jbschlosser and miladm March 30, 2023 22:11

github-actions bot added the ciflow/inductor label Mar 30, 2023

github-actions bot requested a review from SherlockNoMad March 30, 2023 22:11

github-actions bot added the module: dynamo label Mar 30, 2023

github-actions bot requested review from voznesenskym and wconstab March 30, 2023 22:11

bdhirsh added a commit that referenced this pull request Mar 30, 2023

change torch._dynamo.export(aten_graph=...) to allow pre_autograd tra…

029d9bd

…cing ghstack-source-id: 7c13ae9d8f47d4bb85c7ab5b7b49c3cad7ec6fd9 Pull Request resolved: #98031

albanD removed their request for review March 31, 2023 00:11

guangy10 reviewed Mar 31, 2023

View reviewed changes

ezyang reviewed Apr 3, 2023

View reviewed changes

bdhirsh mentioned this pull request Apr 11, 2023

fix scalar-tensor issue with CrossEntropyLoss label smoothing, symintify #98848

Closed

ezyang reviewed Apr 11, 2023

View reviewed changes

bdhirsh mentioned this pull request Apr 12, 2023

AOTAutograd: fix 'Trying to backward through the graph a second time' error #98960

Closed

pytorch-bot bot added the release notes: fx release notes category label Apr 12, 2023

bdhirsh added 3 commits April 12, 2023 19:32

ezyang approved these changes Apr 15, 2023

View reviewed changes

bdhirsh mentioned this pull request Apr 24, 2023

functionalization: error during mutations on mem overlap #99919

Closed

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Apr 24, 2023

pytorchmergebot added the merging label Apr 24, 2023

bdhirsh added 4 commits April 24, 2023 22:29

pytorchmergebot added Merged and removed merging labels Apr 25, 2023

pytorchmergebot closed this in 15e1bee Apr 25, 2023

facebook-github-bot deleted the gh/bdhirsh/404/head branch June 8, 2023 15:47

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

change torch._dynamo.export(aten_graph=...) to allow pre_autograd tracing #98031

change torch._dynamo.export(aten_graph=...) to allow pre_autograd tracing #98031

bdhirsh commented Mar 30, 2023 •

edited

pytorch-bot bot commented Mar 30, 2023 •

edited

guangy10 Mar 31, 2023

guangy10 Mar 31, 2023

ezyang Apr 3, 2023

bdhirsh Apr 3, 2023

ezyang Apr 3, 2023

bdhirsh Apr 11, 2023

bdhirsh commented Apr 11, 2023

ezyang Apr 11, 2023

ezyang Apr 12, 2023

bdhirsh Apr 14, 2023

ezyang Apr 15, 2023

ezyang commented Apr 11, 2023

bdhirsh commented Apr 24, 2023

pytorchmergebot commented Apr 24, 2023

pytorchmergebot commented Apr 24, 2023

bdhirsh commented Apr 25, 2023

pytorchmergebot commented Apr 25, 2023

	"none": export a graph with Python operations. Default is False.
	"none": export a graph with Python operations.

		"aten": export a graph with ATen operations.
		"aten_pre_autograd": export a graph with pre_autograd ATen operations.

change torch._dynamo.export(aten_graph=...) to allow pre_autograd tracing #98031

change torch._dynamo.export(aten_graph=...) to allow pre_autograd tracing #98031

Conversation

bdhirsh commented Mar 30, 2023 • edited

pytorch-bot bot commented Mar 30, 2023 • edited

🔗 Helpful Links

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

✅ No Failures

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

bdhirsh commented Apr 11, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ezyang commented Apr 11, 2023

bdhirsh commented Apr 24, 2023

pytorchmergebot commented Apr 24, 2023

Merge started

pytorchmergebot commented Apr 24, 2023

Merge failed

bdhirsh commented Apr 25, 2023

pytorchmergebot commented Apr 25, 2023

Merge started

bdhirsh commented Mar 30, 2023 •

edited

pytorch-bot bot commented Mar 30, 2023 •

edited