Add functionalization pass in TorchDynamo #99461

tugsbayasgalan · 2023-04-18T21:13:08Z

Stack from ghstack (oldest at bottom):

-> Add functionalization pass in TorchDynamo #99461

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

Differential Revision: D45106409

[ghstack-poisoned]

pytorch-bot · 2023-04-18T21:13:11Z

🔗 Helpful Links

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

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

❌ 5 New Failures

As of commit ae41c99:

NEW FAILURES - The following jobs have failed:

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

anijain2305 · 2023-04-18T21:23:54Z

torch/_dynamo/config.py

 guard_nn_modules = False

+# Optionally functionalize the produced export graph module
+functionalize = False


Lets not add more things to config. Can you use **kwargs in export API for now to do this?

side note - @anijain2305 what's our stance towards BC on torch._dynamo.export? Are we ok with removing kwargs later?

All the options seem equally bad to me. Imo, export is going through a heavy refactoring, and I would wait for aot_export etc to be finalized before we set the API in stone.

anijain2305 · 2023-04-18T21:24:45Z

This looks ok to me for short-term. @bdhirsh Can you take a look?

bdhirsh · 2023-04-18T21:37:46Z

torch/_dynamo/eval_frame.py

+        torch._enable_functionalization(reapply_views=True)
+        yield pytree.tree_map(to_fun, args)
+    finally:
+        torch._disable_functionalization()


I think this will still cause problems, if your Interpreter(graph).run() call lower down raises an exception. You need that to be wrapped by the try/finally to make sure that functionalization is always properly turned off afterwards.

One way would be to change this context manager into a decorator that wraps your interpreter call. Or you can just not bother with the decorator here, and manually try/finally around the interpreter call.

Good catch! I will just do plain try/finally.

bdhirsh · 2023-04-18T21:38:05Z

Otherwise look ok to me!

Fixes: #99000 cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

ghstack-source-id: f3099a6 Pull Request resolved: #99461

bdhirsh

lgtm!

tugsbayasgalan · 2023-04-18T23:24:50Z

@tugsbayasgalan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

Fixes: #99000 cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire Differential Revision: [D45106409](https://our.internmc.facebook.com/intern/diff/D45106409) [ghstack-poisoned]

tugsbayasgalan · 2023-04-18T23:47:22Z

@tugsbayasgalan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

Fixes: #99000 cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire Differential Revision: [D45106409](https://our.internmc.facebook.com/intern/diff/D45106409) [ghstack-poisoned]

ghstack-source-id: ea5b76b Pull Request resolved: #99461

tugsbayasgalan · 2023-04-19T00:20:35Z

@tugsbayasgalan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

Fixes: #99000 cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire Differential Revision: [D45106409](https://our.internmc.facebook.com/intern/diff/D45106409) [ghstack-poisoned]

@soumith

Pull Request resolved: #99461 Fixes: #99000 cc @soumith @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire Differential Revision: [D45106409](https://our.internmc.facebook.com/intern/diff/D45106409/) ghstack-source-id: 187846232

gmagogsfm · 2023-05-03T00:04:07Z

torch/_dynamo/eval_frame.py

+    if functionalize and not aten_graph:
+        raise UserError(
+            UserErrorType.ANTI_PATTERN,
+            "TorchDynamo won't functionalize non-aten graphs. Please set `functionalize` to true",


Suggested change

"TorchDynamo won't functionalize non-aten graphs. Please set `functionalize` to true",

"TorchDynamo won't functionalize non-aten graphs. Please set `aten_graph` to true",

gmagogsfm · 2023-05-03T00:06:16Z

torch/_dynamo/eval_frame.py

            with torch.fx.traceback.preserve_node_meta():
-                return torch.fx.Interpreter(graph).run(*args)
+                if functionalize:
+                    torch._enable_functionalization(reapply_views=False)


I think we concluded to allow views in the graph, should we have reapply_views=True?

Oh lol you are right. We will replace them with view_copy later as pass

gmagogsfm · 2023-05-03T00:06:38Z

torch/_dynamo/eval_frame.py


    new_graph.recompile()

-    # TODO remove this once Executorch uses proper functionalization


Love this deletion!

Fixes: #99000 cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire Differential Revision: [D45106409](https://our.internmc.facebook.com/intern/diff/D45106409) [ghstack-poisoned]

@soumith

Pull Request resolved: #99461 Fixes: #99000 cc @soumith @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire Differential Revision: [D45106409](https://our.internmc.facebook.com/intern/diff/D45106409/) ghstack-source-id: 187851433

tugsbayasgalan · 2023-05-03T22:43:04Z

@pytorchbot rebase

pytorchmergebot · 2023-05-03T22:45:34Z

@pytorchbot successfully started a rebase job. Check the current status here

Fixes: #99000 cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire Differential Revision: [D45106409](https://our.internmc.facebook.com/intern/diff/D45106409) [ghstack-poisoned]

pytorchmergebot · 2023-05-03T22:45:56Z

Successfully rebased gh/tugsbayasgalan/110/orig onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via ghstack checkout https://github.com/pytorch/pytorch/pull/99461)

Fixes: #99000 cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire Differential Revision: [D45106409](https://our.internmc.facebook.com/intern/diff/D45106409) [ghstack-poisoned]

@soumith

Pull Request resolved: #99461 Fixes: #99000 cc @soumith @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire Differential Revision: [D45106409](https://our.internmc.facebook.com/intern/diff/D45106409/) **NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D45106409/)! ghstack-source-id: 188007117

Fixes: #99000 cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire Differential Revision: [D45106409](https://our.internmc.facebook.com/intern/diff/D45106409) [ghstack-poisoned]

@soumith

Pull Request resolved: #99461 Fixes: #99000 cc @soumith @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire Differential Revision: [D45106409](https://our.internmc.facebook.com/intern/diff/D45106409/) **NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D45106409/)! ghstack-source-id: 188187442

facebook-github-bot · 2023-05-05T16:02:53Z

@pytorchbot merge

(Initiating merge automatically since Phabricator Diff has merged)

facebook-github-bot · 2023-05-05T16:05:19Z

@pytorchbot merge

(Initiating merge automatically since Phabricator Diff has merged)

pytorchmergebot · 2023-05-05T16:08:09Z

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

Add functionalization pass in TorchDynamo

2e86c86

[ghstack-poisoned]

github-actions bot added ciflow/inductor module: dynamo labels Apr 18, 2023

tugsbayasgalan requested a review from anijain2305 April 18, 2023 21:19

tugsbayasgalan added the topic: not user facing topic category label Apr 18, 2023

tugsbayasgalan requested a review from bdhirsh April 18, 2023 21:21

anijain2305 reviewed Apr 18, 2023

View reviewed changes

bdhirsh reviewed Apr 18, 2023

View reviewed changes

Update on "Add functionalization pass in TorchDynamo"

fa355e2

Fixes: #99000 cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

Update on "Add functionalization pass in TorchDynamo"

9361b87

Fixes: #99000 cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

tugsbayasgalan added a commit that referenced this pull request Apr 18, 2023

Add functionalization pass in TorchDynamo

5601220

ghstack-source-id: f3099a6 Pull Request resolved: #99461

bdhirsh approved these changes Apr 18, 2023

View reviewed changes

tugsbayasgalan requested review from anijain2305 and bdhirsh April 18, 2023 21:56

bdhirsh approved these changes Apr 18, 2023

View reviewed changes

tugsbayasgalan added a commit that referenced this pull request Apr 19, 2023

Add functionalization pass in TorchDynamo

ab8f2d7

ghstack-source-id: ea5b76b Pull Request resolved: #99461

anijain2305 approved these changes Apr 19, 2023

View reviewed changes

gmagogsfm reviewed May 3, 2023

View reviewed changes

zou3519 mentioned this pull request May 3, 2023

Support map autograd and pytree in/out #100494

Closed

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

pytorchmergebot added the merging label May 5, 2023

pytorchmergebot added Merged and removed merging labels May 5, 2023

pytorchmergebot closed this in bf08b07 May 5, 2023

malfet mentioned this pull request May 17, 2023

[Reland] Add sym_size/stride/numel/storage_offset to native_function.… #100749

Closed

facebook-github-bot deleted the gh/tugsbayasgalan/110/head branch June 8, 2023 18:55

	"TorchDynamo won't functionalize non-aten graphs. Please set `functionalize` to true",
	"TorchDynamo won't functionalize non-aten graphs. Please set `aten_graph` to true",


		new_graph.recompile()

		# TODO remove this once Executorch uses proper functionalization

Add functionalization pass in TorchDynamo #99461

Add functionalization pass in TorchDynamo #99461

Uh oh!

Conversation

tugsbayasgalan commented Apr 18, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Apr 18, 2023 • 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/99461

❌ 5 New Failures

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!

anijain2305 commented Apr 18, 2023

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bdhirsh commented Apr 18, 2023

Uh oh!

bdhirsh left a comment

Choose a reason for hiding this comment

Uh oh!

tugsbayasgalan commented Apr 18, 2023

Uh oh!

tugsbayasgalan commented Apr 18, 2023

Uh oh!

tugsbayasgalan commented Apr 19, 2023

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!

tugsbayasgalan commented May 3, 2023

Uh oh!

pytorchmergebot commented May 3, 2023

Uh oh!

pytorchmergebot commented May 3, 2023

Uh oh!

facebook-github-bot commented May 5, 2023

Uh oh!

facebook-github-bot commented May 5, 2023

Uh oh!

pytorchmergebot commented May 5, 2023

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

tugsbayasgalan commented Apr 18, 2023 •

edited

Loading

pytorch-bot bot commented Apr 18, 2023 •

edited

Loading