Support nn_module_stack in torch.export(strict=False) by tugsbayasgalan · Pull Request #115454 · pytorch/pytorch

tugsbayasgalan · 2023-12-08T21:33:34Z

Stack from ghstack (oldest at bottom):

-> Support nn_module_stack in torch.export(strict=False) #115454

cc @avikchaudhuri @gmagogsfm @zhxchen17 @angelayi @suo @ydwu4

[ghstack-poisoned]

pytorch-bot · 2023-12-08T21:33:39Z

🔗 Helpful Links

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

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

❗ 1 Merge Blocking SEVs

There is 1 active merge blocking SEVs. Please view them below:

(merge blocking) Trunk is very red, let's pause everything except fixing it

If you must merge, use @pytorchbot merge -f.

✅ You can merge normally! (2 Unrelated Failures)

As of commit 874ec75 with merge base 1474eb5 ():

UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:

pull / linux-focal-py3.11-clang10 / test (dynamo, 5, 7, linux.2xlarge, unstable) (gh)
test_unary_ufuncs.py::TestUnaryUfuncsCPU::test_reference_numerics_large__refs_special_multigammaln_mvlgamma_p_5_cpu_int16
pull / linux-focal-py3.8-clang10 / test (dynamo, 5, 7, linux.2xlarge, unstable) (gh)
test_unary_ufuncs.py::TestUnaryUfuncsCPU::test_reference_numerics_large__refs_special_multigammaln_mvlgamma_p_5_cpu_int16

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

ghstack-source-id: 1634a3b Pull Request resolved: #115454

[ghstack-poisoned]

ghstack-source-id: 832bed4 Pull Request resolved: #115454

torch/_functorch/_aot_autograd/schemas.py

suo · 2023-12-09T21:21:18Z

torch/fx/experimental/proxy_tensor.py

+            return super().getattr(attr, attr_val, parameter_proxy_cache)
+        return self.proxy_type(attr_val, attr)
+
+    def call_module(self, m, forward, args, kwargs):


This excludes the preserve_signature handling of the version in aot_export. We need this logic, right? Otherwise preserve_signature does not function correctly for non-strict mode?

I'm okay with taking this as a follow-up, but let's make sure to track it

Yep will do that as follow up.

torch/export/_trace.py

[ghstack-poisoned]

ghstack-source-id: 29cca95 Pull Request resolved: #115454

torch/_higher_order_ops/map.py

torch/fx/experimental/proxy_tensor.py

torch/_functorch/_aot_autograd/schemas.py

test/export/test_export.py

torch/_functorch/_aot_autograd/traced_function_transforms.py

bdhirsh · 2023-12-13T17:28:57Z

torch/fx/experimental/proxy_tensor.py

        torch.set_autocast_cache_enabled(old_value)


+class ModuleStackTracer(PythonKeyTracer):


cc @zou3519 (Ed's out - anyone else you think should review this?)

torch/fx/experimental/proxy_tensor.py

tugsbayasgalan · 2023-12-19T23:58:44Z

@pytorchbot merge

pytorchmergebot · 2023-12-20T00:01:05Z

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

jeanschmidt · 2023-12-21T13:57:01Z

@pytorchbot revert -m "Breaking internal tests recycle_bin_citadel and executorch, check internal diff to see more details" -c nosignal

pytorchmergebot · 2023-12-21T13:58:47Z

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

pytorchmergebot · 2023-12-21T13:58:52Z

Reverting PR 115454 failed

Reason: Command git -C /home/runner/work/pytorch/pytorch revert --no-edit 6730b5bcb41e0519572759d9ad9852a113d0a7e4 returned non-zero exit code 1

Auto-merging test/export/test_export.py
Auto-merging test/test_proxy_tensor.py
Auto-merging torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py
CONFLICT (content): Merge conflict in torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py
Auto-merging torch/_functorch/_aot_autograd/traced_function_transforms.py
Auto-merging torch/_functorch/aot_autograd.py
CONFLICT (content): Merge conflict in torch/_functorch/aot_autograd.py
Auto-merging torch/export/_trace.py
Auto-merging torch/fx/experimental/proxy_tensor.py
error: could not revert 6730b5bcb41... Support nn_module_stack in torch.export(strict=False) (#115454)
hint: After resolving the conflicts, mark them with
hint: "git add/rm <pathspec>", then run
hint: "git revert --continue".
hint: You can instead skip this commit with "git revert --skip".
hint: To abort and get back to the state before "git revert",
hint: run "git revert --abort".

Details for Dev Infra team

Raised by workflow job

Pull Request resolved: #115188 Approved by: https://github.com/bdhirsh

This reverts commit a267d67. Reverted #115188 on behalf of https://github.com/jeanschmidt due to sadly, it is required to revert this commit in order to revert #115454 ([comment](#115188 (comment)))

jeanschmidt · 2023-12-21T14:03:40Z

@pytorchbot revert -m "Breaking internal tests recycle_bin_citadel and executorch, check internal diff to see more details" -c nosignal

pytorchmergebot · 2023-12-21T14:05:39Z

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

pytorchmergebot · 2023-12-21T14:05:48Z

@tugsbayasgalan your PR has been successfully reverted.

This reverts commit 6730b5b. Reverted #115454 on behalf of https://github.com/jeanschmidt due to Breaking internal tests recycle_bin_citadel and executorch, check internal diff to see more details ([comment](#115454 (comment)))

malfet · 2023-12-21T20:05:58Z

cc: @huydhn @mergennachin do you know why wasn't it caught in OSS, as we run executorch builds there now, don't we?

Pull Request resolved: pytorch#115454 Approved by: https://github.com/suo, https://github.com/bdhirsh

Summary: X-link: pytorch/pytorch#115454 Approved by: https://github.com/suo, https://github.com/bdhirsh Differential Revision: D52339616 Pulled By: tugsbayasgalan

tugsbayasgalan · 2023-12-25T01:08:44Z

Landed in: https://github.com/pytorch/pytorch/pull/116309/files

huydhn · 2023-12-27T07:35:11Z

cc: @huydhn @mergennachin do you know why wasn't it caught in OSS, as we run executorch builds there now, don't we?

The build was green https://github.com/pytorch/pytorch/actions/runs/7269905380/job/19808210062 on OSS and we pre-build ExecuTorch in Docker. However, it was built with against a 2-week old commit from what I see https://github.com/pytorch/pytorch/blob/6730b5bcb41e0519572759d9ad9852a113d0a7e4/.ci/docker/ci_commit_pins/executorch.txt. This could explain why OSS CI missed it. ExecuTorch on fbcode, on the other hand, used a newer commit for sure.

This PR #115599 was supposed to land ExecuTorch latest main commit, but I think GitHub had a hiccup here and pytorchbot didn't stamp the approval as usual. I didn't know about this till last Friday when I stamped the approval myself.

It looks ok now in the recent updates:

Having a way to notify oncall about repeated failed attempts to merge an auto commit hash update might be the answer here (stale commit)

TEST

4b43e22

[ghstack-poisoned]

tugsbayasgalan requested review from avikchaudhuri, gmagogsfm and zhxchen17 as code owners December 8, 2023 21:33

pytorch-bot bot added release notes: AO frontend labels Dec 8, 2023

tugsbayasgalan added a commit that referenced this pull request Dec 8, 2023

TEST

085491f

ghstack-source-id: 1634a3b Pull Request resolved: #115454

github-actions bot added ciflow/inductor module: export labels Dec 8, 2023

Update on "TEST"

98fabde

[ghstack-poisoned]

Update on "TEST"

507bd57

[ghstack-poisoned]

tugsbayasgalan added a commit that referenced this pull request Dec 9, 2023

Support nn_module_stack in torch.export(strict=False)

13c8be4

ghstack-source-id: 832bed4 Pull Request resolved: #115454

tugsbayasgalan changed the title ~~TEST~~ [WIP] Support nn_module_stack in torch.export(strict=False) Dec 9, 2023

tugsbayasgalan requested review from bdhirsh and suo December 9, 2023 20:21

suo approved these changes Dec 9, 2023

View reviewed changes

zhxchen17 reviewed Dec 11, 2023

View reviewed changes

torch/export/_trace.py Outdated Show resolved Hide resolved

Update on "[WIP] Support nn_module_stack in torch.export(strict=False)"

27cd88c

[ghstack-poisoned]

tugsbayasgalan added a commit that referenced this pull request Dec 13, 2023

Support nn_module_stack in torch.export(strict=False)

b310151

ghstack-source-id: 29cca95 Pull Request resolved: #115454

tugsbayasgalan commented Dec 13, 2023

View reviewed changes

torch/_higher_order_ops/map.py Show resolved Hide resolved

tugsbayasgalan commented Dec 13, 2023

View reviewed changes

torch/fx/experimental/proxy_tensor.py Outdated Show resolved Hide resolved