[functorch] add functorch functional_call, update tests to test this #89213

samdow · 2022-11-17T14:58:27Z

Stack from ghstack (oldest at bottom):

…l_call [ghstack-poisoned]

pytorch-bot · 2022-11-17T14:58:33Z

🔗 Helpful Links

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

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

✅ No Failures

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

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

…g functional_call" [ghstack-poisoned]

…l_call ghstack-source-id: 6c12b87 Pull Request resolved: #89213

…g functional_call" [ghstack-poisoned]

…l_call ghstack-source-id: ca4bfc8 Pull Request resolved: #89213

…g functional_call" [ghstack-poisoned]

… test this" [ghstack-poisoned]

samdow · 2022-12-16T15:52:16Z

@zou3519 stack now starts here. It's ready for review!

zou3519

Code looks reasonable to me. I had some high-level questions (that I might have asked on the other PRs, sorry if they're redundant):

should the new functional_call API be the same as the nn.stateless.utils.functional_call API, or is the plan that we'll deprecate the nn.stateless.utils.functional_call immediately?
are we exposing the new functional_call in functorch, or should it just be exposed in torch.func?
should we have params_and_buffers_disable_autograd_tracking ? Or should we just warn somewhere in the functional_call docs and show an example with .detach() ?

functorch/__init__.py

torch/_functorch/functional_call.py

test/functorch/test_eager_transforms.py

zou3519 · 2022-12-16T16:12:54Z

torch/_functorch/functional_call.py

+def functional_call(
+    module: 'torch.nn.Module',
+    parameter_and_buffer_dicts: Union[Dict[str, Tensor], Tuple[Dict[str, Tensor]]],
+    args: Union[Any, Tuple],
+    kwargs: Dict[str, Any] = None,
+):


I assume the docs are coming later?

did you want nn.utils.stateless.functional_call to also have this API, or does nn.utils.stateless.functional_call have a different API (and we deprecate it?)

Ahh thanks for the catch I forgot to move the docs over

Re: nn.utils.stateless.functional_call, I think we should deprecate it and ask users to use torch.func instead since they are the same API. I'll add a PR to do that. Also since it was technically BC-breaking to change this since previously parameters_and_buffers must have been a single dictionary

torch/_functorch/functional_call.py

test/functorch/test_eager_transforms.py

torch/_functorch/functional_call.py

… test this" [ghstack-poisoned]

zou3519 · 2022-12-20T16:24:33Z

Btw, torch.func exists now, so let's put functional_call there

… test this" [ghstack-poisoned]

samdow · 2022-12-21T23:36:40Z

Btw, torch.func exists now, so let's put functional_call there

Done! Let me know if there's anything else!

… test this" [ghstack-poisoned]

zou3519

I'm a bit confused about the duplicate keys check, but other than that, this LGTM. I had some minor comments on nits and suggestions to improve the testing

zou3519 · 2022-12-22T20:54:34Z

docs/source/func.api.rst

+.. toctree::
+   :maxdepth: 1
+
+   batch_norm


Does this exist? I don't see it in this PR

Good catch! Moved it over but it still refers to functorch.experimental.replace_all_batch_norm_modules_ I was going to add a PR at the top to fix this but happy to squash it in here if that's preferred

zou3519 · 2022-12-22T20:58:00Z

torch/_functorch/functional_call.py

+            >>> mod = nn.Linear(1, 1)
+            >>> d = {k: v.detach() for k, v in mod.named_parameters()}
+            >>> grad(lambda x: functional_call(mod, d, x), torch.randn((1, 1)))  # doesn't tracks grads for params


Nit: some text explanation of what is going on here, perhaps also with a compute_loss function. "applies the grad transform over the parameters of a model".

https://pytorch.org/functorch/stable/generated/functorch.make_functional.html#functorch.make_functional has a nice end-to-end example that we should follow.

Added end-to-end example and moved this down to be a part of the note on not tracking. Here's what the ending looks like, rendered

torch/_functorch/functional_call.py

test/functorch/test_eager_transforms.py

zou3519 · 2022-12-22T21:57:02Z

torch/_functorch/functional_call.py

+    module: 'torch.nn.Module',
+    parameter_and_buffer_dicts: Union[Dict[str, Tensor], Tuple[Dict[str, Tensor], ...]],
+    args: Union[Any, Tuple],
+    kwargs: Dict[str, Any] = None,


Just to check: Did the weight tying update for stateless.functional_call happen already? If so, should functional_call have an option to toggle the weight tying behavior?

#90477 hasn't been approved yet. I can add it in there if this lands first

Sorry, totally missed that PR

… test this" [ghstack-poisoned]

zou3519

LGTM, some minor nits

docs/source/func.api.rst

zou3519 · 2022-12-28T14:11:59Z

docs/source/func.batch_norm.rst

+What's happening?
+-----------------
+Batch Norm requires in-place updates to running_mean and running_var of the same size as the input.
+Functorch does not support inplace update to a regular tensor that takes in a batched tensor (i.e.


nit: vmap? And link section from the UX limitations doc if possible (there is some sphinx syntax to link it)

I think the UX limitations hasn't been pulled over yet (probably my bad for not looking at a PR yet. Filed #91509 so I don't forget to do this in a follow up)

zou3519 · 2022-12-28T14:13:41Z

docs/source/func.batch_norm.rst

+How to fix
+----------
+All of these options assume that you don't need running stats. If you're using a module this means
+that it's assumed you won't use batch norm in evaluation mode. If you have a use case that involves
+running batch norm with vmap in evaluation mode, please file an issue


Are we not going to recommend replacing the batch norm with another normalization layer? That should be good for completeness so the user knows their options

Added language to 1 + 2 on how to replace with GroupNorm. From Opacus' work, generally replacing the norm layer with a different layer is fragile but the version functorch offers (replace_batch_norm_modules_) turns off tracking should work since it just updates the layer in place instead of trying to replace the whole layer

test/test_stateless.py

… test this" [ghstack-poisoned]

pytorchmergebot · 2023-01-03T16:30:33Z

Rebased gh/samdow/46/orig onto refs/remotes/origin/viable/strict because #88850 was rebased, please pull locally before adding more changes (for example, via ghstack checkout https://github.com/pytorch/pytorch/pull/89213)

facebook-github-bot · 2023-01-04T00:33:24Z

This pull request has been merged in c5e5916.

[functorch] update make_functional tests to also test using functiona…

b3eb608

…l_call [ghstack-poisoned]

This was referenced Nov 17, 2022

[functorch] add new ensembling api, demonstrate in example #88850

Closed

[functorch] rewrite examples that use make_functional to use functional_call #88851

Closed

pytorch-bot bot added the topic: not user facing topic category label Nov 17, 2022

Update on "[functorch] update make_functional tests to also test usin…

ac5047b

…g functional_call" [ghstack-poisoned]

samdow pushed a commit that referenced this pull request Nov 17, 2022

[functorch] update make_functional tests to also test using functiona…

b228381

…l_call ghstack-source-id: 6c12b87 Pull Request resolved: #89213

Update on "[functorch] update make_functional tests to also test usin…

bfb3b2f

…g functional_call" [ghstack-poisoned]

samdow pushed a commit that referenced this pull request Nov 18, 2022

[functorch] update make_functional tests to also test using functiona…

e38254d

…l_call ghstack-source-id: ca4bfc8 Pull Request resolved: #89213

samdow added 2 commits December 15, 2022 15:19

Update on "[functorch] update make_functional tests to also test usin…

9e55461

…g functional_call" [ghstack-poisoned]

Update on "[functorch] update make_functional tests to also test usin…

69dff93

…g functional_call" [ghstack-poisoned]

samdow changed the title ~~[functorch] update make_functional tests to also test using functional_call~~ [functorch] add functorch functional_call, update tests to test this Dec 16, 2022

Update on "[functorch] add functorch functional_call, update tests to…

290b0d3

… test this" [ghstack-poisoned]

samdow requested a review from zou3519 December 16, 2022 15:51

zou3519 reviewed Dec 16, 2022

View reviewed changes

samdow added 2 commits December 19, 2022 11:11

Update on "[functorch] add functorch functional_call, update tests to…

585b88a

… test this" [ghstack-poisoned]

Update on "[functorch] add functorch functional_call, update tests to…

1afb3b8

… test this" [ghstack-poisoned]

samdow added 2 commits December 21, 2022 10:00

Update on "[functorch] add functorch functional_call, update tests to…

6402ac8

… test this" [ghstack-poisoned]

Update on "[functorch] add functorch functional_call, update tests to…

f698770

… test this" [ghstack-poisoned]

samdow mentioned this pull request Dec 21, 2022

[functorch] excise remaining functorch imports from examples #91282

Closed

samdow added 3 commits December 22, 2022 13:02

Update on "[functorch] add functorch functional_call, update tests to…

5555b88

… test this" [ghstack-poisoned]

Update on "[functorch] add functorch functional_call, update tests to…

520471d

… test this" [ghstack-poisoned]

Update on "[functorch] add functorch functional_call, update tests to…

8ae3504

… test this" [ghstack-poisoned]

zou3519 reviewed Dec 22, 2022

View reviewed changes

Update on "[functorch] add functorch functional_call, update tests to…

14f4efb

… test this" [ghstack-poisoned]

samdow mentioned this pull request Dec 27, 2022

[functorch] move batch_norm_replacement to torch.func #91412

Closed

Update on "[functorch] add functorch functional_call, update tests to…

e8c0fb6

… test this" [ghstack-poisoned]

zou3519 approved these changes Dec 28, 2022

View reviewed changes

zou3519 mentioned this pull request Dec 28, 2022

[stateless] add weight tying support #90477

Closed

Update on "[functorch] add functorch functional_call, update tests to…

2021b1f

… test this" [ghstack-poisoned]

samdow mentioned this pull request Dec 29, 2022

[functorch] make batch norm docs point to UX limitations #91509

Open

Update on "[functorch] add functorch functional_call, update tests to…

497a034

… test this" [ghstack-poisoned]

pytorchmergebot closed this in c5e5916 Jan 4, 2023

facebook-github-bot added the Merged label Jan 4, 2023

albanD mentioned this pull request Jan 13, 2023

NN Module functional API #49171

Closed

facebook-github-bot deleted the gh/samdow/48/head branch June 8, 2023 18:43

[functorch] add functorch functional_call, update tests to test this #89213

[functorch] add functorch functional_call, update tests to test this #89213

Uh oh!

Conversation

samdow commented Nov 17, 2022 • edited by pytorchmergebot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Nov 17, 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/89213

✅ No Failures

Uh oh!

samdow commented Dec 16, 2022

Uh oh!

zou3519 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zou3519 commented Dec 20, 2022

Uh oh!

samdow commented Dec 21, 2022

Uh oh!

zou3519 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!

Uh oh!

Uh oh!

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!

zou3519 left a comment

Choose a reason for hiding this comment

Uh oh!

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!

Uh oh!

Uh oh!

pytorchmergebot commented Jan 3, 2023

Uh oh!

facebook-github-bot commented Jan 4, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

samdow commented Nov 17, 2022 •

edited by pytorchmergebot

Loading

pytorch-bot bot commented Nov 17, 2022 •

edited

Loading