[jit] reuse stubs_fn whenever possible to create submodule #43872

wanchaol · 2020-08-31T05:33:35Z

Stack from ghstack:

[jit] reuse stubs_fn whenever possible to create submodule #43872 [jit] reuse stubs_fn whenever possible to create submodule

when compiling submodule, we always use the default method inference rule, which includes forward and methods it calls, this is also the case for tracing. For tracing we don't want to compile the forward by default for submodule. This PR allows the recursive scripting to reuse stubs_fn to create its submodule instead of alwaysing using the default rule. So three cases here:

tracing: only script exported methods, still trace forward and its
recursive calls
scripting: always reuse stubs_fn
legacy ScriptModule: recursive compile use default infer_methods
rule, if the module is a scriptModule, find methods from
self._methods

Fixes #43729

Differential Revision: D23430176

This PR allows the recursive scripting to have a separate submodule_stubs_fn to create its submodule with specific user provided rules. Fixes #43729 [ghstack-poisoned]

This PR allows the recursive scripting to have a separate submodule_stubs_fn to create its submodule with specific user provided rules. Fixes #43729 ghstack-source-id: 040e5b9 Pull Request resolved: #43872

codecov · 2020-08-31T08:22:56Z

Codecov Report

❗ No coverage uploaded for pull request base (gh/wanchaol/125/base@8501b89). Click here to learn what that means.
The diff coverage is n/a.

@@                   Coverage Diff                   @@
##             gh/wanchaol/125/base   #43872   +/-   ##
=======================================================
  Coverage                        ?   67.85%           
=======================================================
  Files                           ?      384           
  Lines                           ?    49917           
  Branches                        ?        0           
=======================================================
  Hits                            ?    33871           
  Misses                          ?    16046           
  Partials                        ?        0

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 8501b89...652f786. Read the comment docs.

This PR allows the recursive scripting to have a separate submodule_stubs_fn to create its submodule with specific user provided rules. Fixes #43729 Differential Revision: [D23430176](https://our.internmc.facebook.com/intern/diff/D23430176) [ghstack-poisoned]

This PR allows the recursive scripting to have a separate submodule_stubs_fn to create its submodule with specific user provided rules. Fixes #43729 ghstack-source-id: 2eb56f9 Pull Request resolved: #43872

dr-ci · 2020-09-01T04:42:20Z

💊 CI failures summary and remediations

As of commit 652f786 (more details on the Dr. CI page):

✅ None of the CI failures appear to be your fault 💚

1/1 broken upstream at merge base 8501b89 on Sep 22 from 10:00am to 11:35am PDT (4 commits; 4b42f0b - ef885c1)

🚧 1 fixed upstream failure:

These were probably caused by upstream breakages that were already fixed.

Please rebase on the viable/strict branch (expand for instructions)

If your commit is newer than viable/strict, you can try basing on an older, stable commit:

git fetch https://github.com/pytorch/pytorch viable/strict
git rebase --onto FETCH_HEAD $(git merge-base origin/master HEAD)

If your commit is older than viable/strict:

git fetch https://github.com/pytorch/pytorch viable/strict
git rebase FETCH_HEAD

Check out the recency history of this "viable master" tracking branch.

pytorch_xla_linux_bionic_py3_6_clang9_build on Sep 22 from 10:00am to 11:35am PDT (4 commits; 4b42f0b - ef885c1)
- 🔁 rerun

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group.

See how this bot performed.

This comment has been revised 25 times.

eellison · 2020-09-02T19:27:00Z

torch/jit/_recursive.py

    concrete_type._create_methods(defs, rcbs, defaults)

-def create_script_module(nn_module, stubs_fn, share_types=True):
+def create_script_module(nn_module, stubs_fn, submodule_stubs_fn=None, share_types=True):


Do we need to add this extra argument? Is there ever going to be a case where we don't use all the functions that are marked with _jit_internal.FunctionModifiers.EXPORT ?

Before this PR, the rule is: submodule will be compiled using the default inference rule, which only contains the forward method + export methods, so export methods will always be compiled before/after this PR. This argument is added because we don't want tracing to compile the forward method, we want the tracing to compile the export methods, but still trace the forward calls

Isn't this what stubs_fn is already for ?

stubs_fn is only used for the current module compilation, when we are doing recursive compilation for the current module's submodule, we are using infer_methods_to_compile as the stubs_fn, which disallow us from compiling the submodule with different rules. I was trying to also use the stubs_fn as the rule for submodule compilation to replace infer_methods_to_compile, but it failed because we need to maintain the legacy ScriptModule api where submodule don't have cls._methods available. See

pytorch/torch/jit/_script.py

Line 203 in 7816d53

] = torch.jit._recursive.create_script_module(self, make_stubs)

Copied from chat:

we are creating ScriptModule using make_stubs function, and the current ScriptModule initialization use cls._methods as stubs_fn infer rule, but that ScriptModule's submodule is compiled using recursive compilation, if we just use stubs_fn as the submodule method infer rule, the submodule compilation also needs cls._methods, which is not available as cls._methods only available in the current ScriptModule (we construct cls._methods in ScriptMeta.__init__

eellison

PR looks good, trying to understand current code organization

yf225 · 2020-09-08T21:39:32Z

@wanchaol @eellison Thanks so much for you guys' help on this! After this is landed we can implement the replace_submodule_with_scripted_version API (D23266907), which is a blocker for launching QPS tests for mixed tracing+scripting models for PyPer :)

suo · 2020-09-09T23:36:58Z

Can we clarify the context here? I see that this is about something involving mixing tracing and scripting, but what specific use case are we trying to make work? The added test case covers functionality that should already exist today (when tracing a module, we should attempt to compile any @exported methods), so it's not very illuminating.

wanchaol · 2020-09-10T00:31:26Z

Can we clarify the context here? I see that this is about something involving mixing tracing and scripting, but what specific use case are we trying to make work? The added test case covers functionality that should already exist today (when tracing a module, we should attempt to compile any @exported methods), so it's not very illuminating.

@suo commented with a detailed code snippet in the issue #43729 (comment)

torch/jit/_recursive.py

This PR allows the recursive scripting to have a separate submodule_stubs_fn to create its submodule with specific user provided rules. Fixes #43729 Differential Revision: [D23430176](https://our.internmc.facebook.com/intern/diff/D23430176) [ghstack-poisoned]

This PR allows the recursive scripting to have a separate submodule_stubs_fn to create its submodule with specific user provided rules. Fixes #43729 ghstack-source-id: 455e115 Pull Request resolved: #43872

This PR allows the recursive scripting to have a separate submodule_stubs_fn to create its submodule with specific user provided rules. Fixes #43729 Differential Revision: [D23430176](https://our.internmc.facebook.com/intern/diff/D23430176) [ghstack-poisoned]

… allow reuse rule for submodule" This PR allows the recursive scripting to reuse stubs_fn to create its submodule. Fixes #43729 Differential Revision: [D23430176](https://our.internmc.facebook.com/intern/diff/D23430176) [ghstack-poisoned]

…e rule for submodule This PR allows the recursive scripting to reuse stubs_fn to create its submodule. Fixes #43729 ghstack-source-id: 5c56845 Pull Request resolved: #43872

… allow reuse rule for submodule" when compiling submodule, we always use the default method inference rule, which includes `forward` and methods it calls, this is also the case for tracing. For tracing we don't want to compile the forward by default for submodule. This PR allows the recursive scripting to reuse stubs_fn to create its submodule instead of alwaysing using the default rule. Fixes #43729 Differential Revision: [D23430176](https://our.internmc.facebook.com/intern/diff/D23430176) [ghstack-poisoned]

…e rule for submodule when compiling submodule, we always use the default method inference rule, which includes `forward` and methods it calls, this is also the case for tracing. For tracing we don't want to compile the forward by default for submodule. This PR allows the recursive scripting to reuse stubs_fn to create its submodule instead of alwaysing using the default rule. Fixes #43729 ghstack-source-id: 2201f7b Pull Request resolved: #43872

suo · 2020-09-18T18:23:07Z

torch/jit/_recursive.py

    return concrete_type

-def create_script_module(nn_module, stubs_fn, share_types=True):
+def create_script_module(nn_module, stubs_fn, share_types=True, reuse_stubs_fn=False):


As far as I can tell, reuse_stubs_fn is always True right? Because either:

You are tracing, and you want to keep tracing the forwards of the other modules

You are scripting, and you are using the default stubs_fn.

Either way, the stubs_fn is the same within a given recursive run, right? If that's the case, can we just remove this option and ensure that we always re-use the same stubs fn?

While this is almost true and it was the methodology that I was trying, but unfortunately it fails, because we have one exception: the Legacy ScriptModule compilation, if you look at here

pytorch/torch/jit/_script.py

Line 203 in 7816d53

] = torch.jit._recursive.create_script_module(self, make_stubs)

we are using the make_stubs rule to compile the legacy ScriptModule as entry point, but inside the recursive run, we are using the default infer_methods_to_compile rule for submodule compilation. I think this is because cls._methods.items() is not available in ScriptModule's submodule bc that might be a plain nn.Module and don't have cls._methods crafted upon construction.

In fact, if I always use the same stubs_fn, it will crash on some tests and complain something like this:

File "/home/wanchaol/pytorch/torch/jit/_recursive.py", line 364, in create_script_module_impl method_stubs = stubs_fn(nn_module) File "/home/wanchaol/pytorch/torch/jit/_script.py", line 202, in make_stubs return [v for k, v in sorted(cls._methods.items())] AttributeError: type object 'Over' has no attribute '_methods'

I think we should either: 1. use the approach I proposed here, so that we don't reuse_stubs_fn in legacyScriptModule 2. disable recursive compilation in legacy ScriptModule API, only allow user to construct their submodule with all modules inheriting from torch.jit.ScriptModule (or say, torch.jit.ScriptModule should only contain submodule with inheriting from torch.jit.ScriptModule.)

I would prefer the first option bc it does not BC break the behavior that we already exposed to the user.

In general, I would prefer to add more to our legacy path shims, rather than change the current behavior just because the legacy path has some constraints. That way we can keep the current behavior clean and add the complexity in the legacy path.

In this case, that would mean: change the custom make_stubs used in the the legacy ScriptModule compilation to fork its behavior, so it would look like:

def make_stubs: if I am compiled the legacy way: look in _methods to make stubs else: use the regular `infer_methods_to_compile`

when compiling submodule, we always use the default method inference rule, which includes `forward` and methods it calls, this is also the case for tracing. For tracing we don't want to compile the forward by default for submodule. This PR allows the recursive scripting to reuse stubs_fn to create its submodule instead of alwaysing using the default rule. So three cases here: 1. tracing: only script exported methods, still trace forward and its recursive calls 2. scripting: always reuse stubs_fn 3. legacy ScriptModule: recursive compile use default infer_methods rule, if the module is a scriptModule, find methods from self._methods Fixes #43729 Differential Revision: [D23430176](https://our.internmc.facebook.com/intern/diff/D23430176) [ghstack-poisoned]

when compiling submodule, we always use the default method inference rule, which includes `forward` and methods it calls, this is also the case for tracing. For tracing we don't want to compile the forward by default for submodule. This PR allows the recursive scripting to reuse stubs_fn to create its submodule instead of alwaysing using the default rule. So three cases here: 1. tracing: only script exported methods, still trace forward and its recursive calls 2. scripting: always reuse stubs_fn 3. legacy ScriptModule: recursive compile use default infer_methods rule, if the module is a scriptModule, find methods from self._methods Fixes #43729 ghstack-source-id: 73f52f8 Pull Request resolved: #43872

suo

lgtm!

facebook-github-bot · 2020-09-23T22:19:13Z

@wanchaol merged this pull request in 3f89b77.

[jit] allow submodule methods inference rule be different

f0e4fcd

This PR allows the recursive scripting to have a separate submodule_stubs_fn to create its submodule with specific user provided rules. Fixes #43729 [ghstack-poisoned]

wanchaol requested a review from apaszke as a code owner August 31, 2020 05:33

facebook-github-bot added the oncall: jit Add this issue/PR to JIT oncall triage queue label Aug 31, 2020

wanchaol requested review from eellison and suo August 31, 2020 17:27

eellison reviewed Sep 2, 2020

View reviewed changes

wanchaol mentioned this pull request Sep 10, 2020

[jit] allow tracing to only compile export methods, instead of generating scriptmodule for submodule #43729

Closed

wanchaol requested a review from gmagogsfm September 10, 2020 23:13

gmagogsfm reviewed Sep 11, 2020

View reviewed changes

torch/jit/_recursive.py Outdated Show resolved Hide resolved

wanchaol changed the title ~~[jit] allow submodule methods inference rule be different~~ [jit] add reuse_stubs_fn option to create_script_module to allow reuse rule for submodule Sep 17, 2020

wanchaol requested review from eellison and gmagogsfm September 17, 2020 17:33

suo reviewed Sep 18, 2020

View reviewed changes

wanchaol requested a review from suo September 21, 2020 16:43

wanchaol changed the title ~~[jit] add reuse_stubs_fn option to create_script_module to allow reuse rule for submodule~~ [jit] reuse stubs_fn whenever possible to create submodule Sep 22, 2020

suo approved these changes Sep 23, 2020

View reviewed changes

facebook-github-bot closed this in 3f89b77 Sep 23, 2020

facebook-github-bot added the merged label Sep 23, 2020

facebook-github-bot deleted the gh/wanchaol/125/head branch September 27, 2020 14:15

mruberry added the Merged label Oct 28, 2020

[jit] reuse stubs_fn whenever possible to create submodule #43872

[jit] reuse stubs_fn whenever possible to create submodule #43872

Uh oh!

Conversation

wanchaol commented Aug 31, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Aug 31, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

dr-ci bot commented Sep 1, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

💊 CI failures summary and remediations

🚧 1 fixed upstream failure:

Uh oh!

eellison Sep 2, 2020

Choose a reason for hiding this comment

Uh oh!

wanchaol Sep 2, 2020

Choose a reason for hiding this comment

Uh oh!

eellison Sep 8, 2020

Choose a reason for hiding this comment

Uh oh!

wanchaol Sep 8, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

eellison Sep 9, 2020

Choose a reason for hiding this comment

Uh oh!

eellison left a comment

Choose a reason for hiding this comment

Uh oh!

yf225 commented Sep 8, 2020

Uh oh!

suo commented Sep 9, 2020

Uh oh!

wanchaol commented Sep 10, 2020

Uh oh!

Uh oh!

suo Sep 18, 2020

Choose a reason for hiding this comment

Uh oh!

wanchaol Sep 18, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

suo Sep 21, 2020

Choose a reason for hiding this comment

Uh oh!

suo left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented Sep 23, 2020

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

wanchaol commented Aug 31, 2020 •

edited

Loading

codecov bot commented Aug 31, 2020 •

edited

Loading

dr-ci bot commented Sep 1, 2020 •

edited

Loading

wanchaol Sep 8, 2020 •

edited

Loading

wanchaol Sep 18, 2020 •

edited

Loading