Refactor how AOTAutograd backends are defined by ezyang · Pull Request #89736 · pytorch/pytorch

ezyang · 2022-11-28T04:08:20Z

Stack from ghstack (oldest at bottom):

There was a lot of strangeness in how AOTAutograd backends were previously defined. This refactor replaces the strangeness with something simple and straightforward. The improvements:

There is no longer a footgun aot_autograd "backend" which doesn't actually work. No more mistyping torch._dynamo.optimize("aot_autograd") when you meant "aot_eager"
Deleted aot_print because it's annoying and anyway there's no uses of it
Instead of having BOTH the backend Subgraph and AotAutogradStrategy, there is now only an aot_autograd function which takes the kwargs to configure AOTAutograd, and then gives you a compiler function that does AOTAutograd given those kwargs. Easy.
The primary downside is that we are now eagerly populating all of the kwargs, and that can get us into import cycle shenanigans. Some cycles I resolved directly (e.g., we now no longer manually disable the forward function before passing it to aot_autograd; aot_autograd it does it for us), but for getting inductor decompositions I had to make it take a lambda so I could lazily populate the decomps later.

New code is 130 lines shorter!

Signed-off-by: Edward Z. Yang ezyang@fb.com

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

Signed-off-by: Edward Z. Yang <ezyang@fb.com> [ghstack-poisoned]

pytorch-bot · 2022-11-28T04:08:22Z

🔗 Helpful Links

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

📄 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 our office hours

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

❌ 2 Failures

As of commit c8c363b:

The following jobs have failed:

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

Signed-off-by: Edward Z. Yang <ezyangfb.com> cc mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen chunyuan-w XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]

There was a lot of strangeness in how AOTAutograd backends were previously defined. This refactor replaces the strangeness with something simple and straightforward. The improvements: - There is no longer a footgun aot_autograd "backend" which doesn't actually work. No more mistyping `torch._dynamo.optimize("aot_autograd")` when you meant "aot_eager" - Deleted aot_print because it's annoying and anyway there's no uses of it - Instead of having BOTH the backend Subgraph and AotAutogradStrategy, there is now only an aot_autograd function which takes the kwargs to configure AOTAutograd, and then gives you a compiler function that does AOTAutograd given those kwargs. Easy. - The primary downside is that we are now eagerly populating all of the kwargs, and that can get us into import cycle shenanigans. Some cycles I resolved directly (e.g., we now no longer manually disable the forward function before passing it to aot_autograd; aot_autograd it does it for us), but for getting inductor decompositions I had to make it take a lambda so I could lazily populate the decomps later. New code is 130 lines shorter! Signed-off-by: Edward Z. Yang <ezyangfb.com> cc mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen chunyuan-w XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 desertfire [ghstack-poisoned]

anjali411 · 2022-11-28T13:43:46Z

functorch/_src/aot_autograd.py

        fw_metadata = _fw_metadata

        @staticmethod
-        @disable_torchdynamo


why were we disabling dynamo earlier and not anymore?

I moved the disable. Previously, we manually disabled dynamo on the forwards function we return. Now, dynamo is responsible for disabling itself on the returned compiled function. See https://github.com/pytorch/pytorch/pull/89736/files#diff-0b094ea719a9acd16c316b7ec6975f6c9825398952e9b3a67f180d20ac592d47R82

I wasn't really planning to do this change but import cycles were annoying

anjali411 · 2022-11-28T13:47:04Z

torch/_dynamo/optimizations/training.py

+
+        force_compile_tiny_graphs = kwargs.pop("force_compile_tiny_graphs", False)
+
+        if count_calls(gm.graph) < 2 and not force_compile_tiny_graphs:


as per the comment below for decomps, should we always force compilation?

Forcing compilation on single op graphs for aot eager truly is pointless. Maybe we should do it anyway for debug purposes? Not sure.

hmm yeah. I like the idea to force recompilation in debug mode anyway.

@Chillee you cool with this idea?

Yeah I'm fine.

is it truly pointless? what if the single op decomposes into some ops that can be fused?

we gonna remove this conditional!

albanD

Sounds good!

albanD · 2022-11-28T14:26:29Z

torch/_dynamo/optimizations/training.py

-
-
-aot_ts = AotTorchscript.compile_fn
+DEBUG = False


Should this be in the config?

Yes, probably. Need to negotiate a name for it.

albanD · 2022-11-28T14:28:57Z

torch/_dynamo/optimizations/training.py

+            kwargs["decompositions"] = kwargs["decompositions"]()
+
+        # TODO: stop monkeypatching here (without even cleaning up, UGH!)
+        functorch.compile.config.use_functionalize = True


Shouldn't we just update the config's default at this point?

Yes, see #89663

There was a lot of strangeness in how AOTAutograd backends were previously defined. This refactor replaces the strangeness with something simple and straightforward. The improvements: - There is no longer a footgun aot_autograd "backend" which doesn't actually work. No more mistyping `torch._dynamo.optimize("aot_autograd")` when you meant "aot_eager" - Deleted aot_print because it's annoying and anyway there's no uses of it - Instead of having BOTH the backend Subgraph and AotAutogradStrategy, there is now only an aot_autograd function which takes the kwargs to configure AOTAutograd, and then gives you a compiler function that does AOTAutograd given those kwargs. Easy. - The primary downside is that we are now eagerly populating all of the kwargs, and that can get us into import cycle shenanigans. Some cycles I resolved directly (e.g., we now no longer manually disable the forward function before passing it to aot_autograd; aot_autograd it does it for us), but for getting inductor decompositions I had to make it take a lambda so I could lazily populate the decomps later. New code is 130 lines shorter! Signed-off-by: Edward Z. Yang <ezyangfb.com> cc mlazos soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen chunyuan-w XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 desertfire [ghstack-poisoned]

There was a lot of strangeness in how AOTAutograd backends were previously defined. This refactor replaces the strangeness with something simple and straightforward. The improvements: - There is no longer a footgun aot_autograd "backend" which doesn't actually work. No more mistyping `torch._dynamo.optimize("aot_autograd")` when you meant "aot_eager" - Deleted aot_print because it's annoying and anyway there's no uses of it - Instead of having BOTH the backend Subgraph and AotAutogradStrategy, there is now only an aot_autograd function which takes the kwargs to configure AOTAutograd, and then gives you a compiler function that does AOTAutograd given those kwargs. Easy. - The primary downside is that we are now eagerly populating all of the kwargs, and that can get us into import cycle shenanigans. Some cycles I resolved directly (e.g., we now no longer manually disable the forward function before passing it to aot_autograd; aot_autograd it does it for us), but for getting inductor decompositions I had to make it take a lambda so I could lazily populate the decomps later. New code is 130 lines shorter! Signed-off-by: Edward Z. Yang <ezyang@fb.com> Pull Request resolved: pytorch#89736 Approved by: https://github.com/anjali411, https://github.com/albanD

Refactor how AOTAutograd backends are defined

448d725

Signed-off-by: Edward Z. Yang <ezyang@fb.com> [ghstack-poisoned]

ezyang mentioned this pull request Nov 28, 2022

Add mypy checking for a few files in torch/_dynamo #89731

Closed

github-actions bot requested review from Chillee, SherlockNoMad, albanD, anjali411, antoniojkim, bdhirsh, miladm, voznesenskym and wconstab November 28, 2022 04:08

github-actions bot added ciflow/inductor module: dynamo labels Nov 28, 2022

github-actions bot added the module: inductor label Nov 28, 2022

ezyang added topic: not user facing topic category release notes: dynamo labels Nov 28, 2022

This was referenced Nov 28, 2022

Change aot_module_simplified to take take arguments directly #89669

Closed

Make aot_module_simplified accept fake tensors #89670

Closed

Use isinstance test rather than exact type test for wrap to fake #89671

Closed

anjali411 reviewed Nov 28, 2022

View reviewed changes

ezyang mentioned this pull request Nov 28, 2022

[UPDATED PROTOTYPE] Use dynamo fake tensor mode in aot_autograd, move aot_autograd compilation to lowering time #89672

Closed

anjali411 reviewed Nov 28, 2022

View reviewed changes

anjali411 approved these changes Nov 28, 2022

View reviewed changes

albanD approved these changes Nov 28, 2022

View reviewed changes

voznesenskym mentioned this pull request Nov 28, 2022

Get CI passing #89773

Closed

pytorchmergebot closed this in b589e72 Nov 28, 2022

facebook-github-bot deleted the gh/ezyang/1602/head branch June 8, 2023 16:35


		force_compile_tiny_graphs = kwargs.pop("force_compile_tiny_graphs", False)

		if count_calls(gm.graph) < 2 and not force_compile_tiny_graphs:



		aot_ts = AotTorchscript.compile_fn
		DEBUG = False

Conversation

ezyang commented Nov 28, 2022 • edited by voznesenskym Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Nov 28, 2022 • 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/89736

❌ 2 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!

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!

albanD left a comment

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

ezyang commented Nov 28, 2022 •

edited by voznesenskym

Loading

pytorch-bot bot commented Nov 28, 2022 •

edited

Loading