[reland 2] Call jit decomp in VariableType to improve forward AD coverage #84976

soulitzer · 2022-09-14T01:03:06Z

Reland of #84675

pytorch-bot · 2022-09-14T01:03:10Z

🔗 Helpful Links

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

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

✅ No Failures, 1 Pending

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

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

facebook-github-bot · 2022-09-14T01:05:09Z

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

facebook-github-bot · 2022-09-14T02:28:41Z

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

facebook-github-bot · 2022-09-14T05:37:30Z

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

…rward AD coverage (pytorch#84151)"" This reverts commit acb4a09. ghstack-source-id: 387eda8 Pull Request resolved: pytorch#84675

facebook-github-bot · 2022-09-14T06:47:28Z

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

facebook-github-bot · 2022-09-14T06:51:32Z

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

facebook-github-bot · 2022-09-14T06:51:41Z

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

facebook-github-bot · 2022-09-14T19:13:14Z

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

soulitzer · 2022-09-14T20:17:37Z

Hopefully internal builds should passing now. The issue was that certain builds (e.g. lite_trainer) depended on the VariableType files, but do not depend on the jit. Including all of those additional dependencies to get jit to work on lite trainer will blow up the binary size so I don't think we want that.

Instead, what I'm doing in this updated PR (see latest commit for new changes) is the same thing that is done with VariableHooksInterface, i.e. provide the apis conditionally depending on whether the library depends on jit.

torch/csrc/autograd/VariableTypeUtilsDependOnOps.h

torch/csrc/jit_decomp_interface.cpp

torch/csrc/jit_decomp_interface.h

zou3519

The code and approach looks good to me. Two main comments:

There's a question of what we want to do in the long term. Should we take the TorchScript dependency, or should we just invoke the decomposition in Python? Since the TorchScript route is easy to do and proven to work, we should go with this PR for the short-to-medium term.
It took me a while to figure out what JitDecompInterface was and why we needed it. We should put a comment with a NOTE somewhere so that folks coming along to read it aren't confused.

torch/csrc/jit/runtime/decomposition_registry.cpp

torch/csrc/jit_decomp_interface.h

soulitzer · 2022-09-15T17:35:05Z

I've confirmed that internal tests are passing now. I've cleaned up this PR a little bit and addressed the comments, so it should be ready for another look!

There's a question of what we want to do in the long term. Should we take the TorchScript dependency, or should we just invoke the decomposition in Python? Since the TorchScript route is easy to do and proven to work, we should go with this PR for the short-to-medium term.

This route makes sense to me as well

zou3519

LGTM, minor comments

torch/csrc/autograd/VariableTypeUtils.h

torch/csrc/jit_decomp_interface.cpp

torch/csrc/jit_decomp_interface.h

facebook-github-bot · 2022-09-15T18:28:59Z

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

soulitzer · 2022-09-15T20:18:28Z

@pytorchbot merge -g

pytorchmergebot · 2022-09-15T20:19:48Z

@pytorchbot successfully started a merge job. Check the current status here.
The merge job was triggered with the green (-g) flag. This means that your change will be merged once all checks on your PR have passed (ETA: 0-4 Hours). If this is not the intended behavior, feel free to use some of the other merge options in the wiki.
Please reach out to the PyTorch DevX Team with feedback or questions!

github-actions · 2022-09-15T22:47:05Z

Hey @soulitzer.
You've committed this PR, but it does not have both a 'release notes: ...' and 'topics: ...' label. Please add one of each to the PR. The 'release notes: ...' label should represent the part of PyTorch that this PR changes (fx, autograd, distributed, etc) and the 'topics: ...' label should represent the kind of PR it is (not user facing, new feature, bug fix, perf improvement, etc). The list of valid labels can be found here for the 'release notes: ...' and here for the 'topics: ...'.
For changes that are 'topic: not user facing' there is no need for a release notes label.

…rage (#84976) Reland of #84675 Pull Request resolved: #84976 Approved by: https://github.com/zou3519

soulitzer requested review from mruberry, ngimel and albanD as code owners September 14, 2022 01:03

pytorch-bot bot added the release notes: jit release notes category label Sep 14, 2022

facebook-github-bot added the cla signed label Sep 14, 2022

facebook-github-bot added oncall: jit Add this issue/PR to JIT oncall triage queue fx labels Sep 14, 2022

soulitzer force-pushed the jvp-decomp branch from 8138c1e to 5b87146 Compare September 14, 2022 02:27

soulitzer force-pushed the jvp-decomp branch from 015068a to 1906c1f Compare September 14, 2022 05:35

Revert "Revert "Call jit decomposition in VariableType to increase fo…

4ca594e

…rward AD coverage (pytorch#84151)"" This reverts commit acb4a09. ghstack-source-id: 387eda8 Pull Request resolved: pytorch#84675

soulitzer force-pushed the jvp-decomp branch from 1906c1f to 7c06046 Compare September 14, 2022 06:46

soulitzer force-pushed the jvp-decomp branch from 7c06046 to 0ab088b Compare September 14, 2022 06:51

zou3519 mentioned this pull request Sep 14, 2022

Remove monkey-patching from functorch #85009

Closed

Access JIT decomposition utils through an interface

67e4113

soulitzer force-pushed the jvp-decomp branch from 0ab088b to 67e4113 Compare September 14, 2022 19:12

soulitzer requested a review from zou3519 September 14, 2022 20:11

soulitzer added ciflow/trunk Trigger trunk jobs on your pull request ciflow/periodic Trigger jobs ran periodically on master (periodic.yml) on the PR labels Sep 14, 2022

soulitzer commented Sep 14, 2022

View reviewed changes

torch/csrc/autograd/VariableTypeUtilsDependOnOps.h Outdated Show resolved Hide resolved

soulitzer commented Sep 14, 2022

View reviewed changes

torch/csrc/jit_decomp_interface.cpp Outdated Show resolved Hide resolved

soulitzer commented Sep 14, 2022

View reviewed changes

torch/csrc/jit_decomp_interface.h Outdated Show resolved Hide resolved

zou3519 reviewed Sep 15, 2022

View reviewed changes

torch/csrc/jit/runtime/decomposition_registry.cpp Outdated Show resolved Hide resolved

torch/csrc/jit_decomp_interface.h Outdated Show resolved Hide resolved

soulitzer force-pushed the jvp-decomp branch 2 times, most recently from c292e2f to 32a2fbb Compare September 15, 2022 17:26

Clean up and address comments

c29d8c2

soulitzer force-pushed the jvp-decomp branch from 32a2fbb to c29d8c2 Compare September 15, 2022 17:28

soulitzer requested a review from zou3519 September 15, 2022 17:35

zou3519 approved these changes Sep 15, 2022

View reviewed changes

torch/csrc/autograd/VariableTypeUtils.h Outdated Show resolved Hide resolved

torch/csrc/autograd/VariableTypeUtils.h Outdated Show resolved Hide resolved

torch/csrc/jit_decomp_interface.cpp Outdated Show resolved Hide resolved

torch/csrc/jit_decomp_interface.h Outdated Show resolved Hide resolved

soulitzer force-pushed the jvp-decomp branch from 80bac63 to 3c59876 Compare September 15, 2022 18:25

Address comments

31dc0c0

soulitzer force-pushed the jvp-decomp branch from 3c59876 to 31dc0c0 Compare September 15, 2022 18:27

pytorchmergebot added the Merged label Sep 15, 2022

pytorchmergebot closed this in 7f88934 Sep 15, 2022

soulitzer added release notes: autograd release notes category topic: improvements topic category and removed release notes: jit release notes category labels Sep 15, 2022

d4l3k mentioned this pull request Sep 16, 2022

multipy/runtime: mock _decomp pytorch/multipy#178

Open

malfet mentioned this pull request Sep 22, 2022

import torch fails on Python-3.11 runtime #85506

Closed

zou3519 mentioned this pull request Sep 22, 2022

Measure impact of JIT decompositions, reconsider the design #85513

Open

mehtanirav pushed a commit that referenced this pull request Oct 4, 2022

[reland 2] Call jit decomp in VariableType to improve forward AD cove…

c09b5e6

…rage (#84976) Reland of #84675 Pull Request resolved: #84976 Approved by: https://github.com/zou3519

[reland 2] Call jit decomp in VariableType to improve forward AD coverage #84976

[reland 2] Call jit decomp in VariableType to improve forward AD coverage #84976

Uh oh!

Conversation

soulitzer commented Sep 14, 2022

Uh oh!

pytorch-bot bot commented Sep 14, 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/84976

✅ No Failures, 1 Pending

Uh oh!

facebook-github-bot commented Sep 14, 2022

Uh oh!

facebook-github-bot commented Sep 14, 2022

Uh oh!

facebook-github-bot commented Sep 14, 2022

Uh oh!

facebook-github-bot commented Sep 14, 2022

Uh oh!

facebook-github-bot commented Sep 14, 2022

Uh oh!

facebook-github-bot commented Sep 14, 2022

Uh oh!

facebook-github-bot commented Sep 14, 2022

Uh oh!

soulitzer commented Sep 14, 2022

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zou3519 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

soulitzer commented Sep 15, 2022

Uh oh!

zou3519 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

facebook-github-bot commented Sep 15, 2022

Uh oh!

soulitzer commented Sep 15, 2022

Uh oh!

pytorchmergebot commented Sep 15, 2022

Uh oh!

github-actions bot commented Sep 15, 2022

Uh oh!

Uh oh!

pytorch-bot bot commented Sep 14, 2022 •

edited

Loading