Fix embedding jvp support by making embedding_renorm ignore forward mode AD #78560

samdow · 2022-05-31T19:16:41Z

On functorch, we started seeing embedding forward mode fail. From looking at it, we figured out that recently embedding got forward mode support enabled and then doing forward mode with embedding and max_norm doesn't work with gradcheck, so it's not checked.

What was happening is that embedding_renorm was setting torch.no_grad() which only turns off the backwards mode AD so functorch's jvp tests were still using forward mode AD during the embedding_renorm call. This makes it so that we don't use forward mode during the embedding_renorm call

facebook-github-bot · 2022-05-31T19:16:48Z

🔗 Helpful links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/78560
📄 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

✅ No Failures (0 Pending)

As of commit 25309c2 (more details on the Dr. CI page):

Expand to see more

💚 💚 Looks good so far! There are no failures yet. 💚 💚

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

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

torch/nn/functional.py

jbschlosser · 2022-05-31T20:28:10Z

torch/nn/functional.py

    with torch.no_grad():
-        torch.embedding_renorm_(weight, input, max_norm, norm_type)
+        set_fwd_grad_enabled(False)


Has a context manager for enabling / disabling fwd grad been discussed before?

Sorry for the confusion, Alban and I talked offline and this is a functorch-only function that I'm guessing worked locally through some dynamic linking or something?

This will change in a moment to use detach instead, though it might still be nice to have an equivalent for forward mode

samdow · 2022-05-31T22:16:41Z

Sorry again about the confusion, PR should be in final state of using detach and having tests. Locally, it fixes functorch and passes these tests (fwd ad one did fail before change to torch)

soulitzer

Thanks!

albanD

Thanks for the update!

samdow · 2022-06-02T13:39:04Z

@pytorchbot merge

github-actions · 2022-06-02T13:41:00Z

Hey @samdow.
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.

malfet · 2022-06-02T17:37:06Z

@pytorchbot revert his broke XLA (on CI and trunk), see https://hud.pytorch.org/pytorch/pytorch/commit/ce7c7bb2a9b21c8a05eedc6b43b277654064a1c9

pytorch-bot · 2022-06-02T17:37:08Z

❌ 🤖 pytorchbot command failed:

@pytorchbot revert: error: the following arguments are required: -m/--message

usage: @pytorchbot revert -m MESSAGE
                          [-c {nosignal,ignoredsignal,landrace,weird,ghfirst}]

```.
Try `@pytorchbot help` for more info.

malfet · 2022-06-02T17:38:11Z

@pytorchmergebot revert -m "broke XLA (on CI and trunk), see https://hud.pytorch.org/pytorch/pytorch/commit/ce7c7bb2a9b21c8a05eedc6b43b277654064a1c9" -c ignoredsignal

malfet · 2022-06-02T17:38:47Z

@pytorchbot revert -m "broke XLA (on CI and trunk), see https://hud.pytorch.org/pytorch/pytorch/commit/ce7c7bb2a9b21c8a05eedc6b43b277654064a1c9" -c ignoredsignal

…orward mode AD (#78560)" This reverts commit ce7c7bb. Reverted #78560 on behalf of https://github.com/malfet due to broke XLA (on CI and trunk), see https://hud.pytorch.org/pytorch/pytorch/commit/ce7c7bb2a9b21c8a05eedc6b43b277654064a1c9

…ode AD (#78560) (#78560) Summary: On functorch, we started seeing [embedding forward mode fail](pytorch/functorch#816). From looking at it, we figured out that recently [embedding got forward mode support enabled](369d9f4) and then doing forward mode with embedding and [max_norm doesn't work with gradcheck](https://github.com/pytorch/pytorch/blob/master/torch/testing/_internal/common_methods_invocations.py#L8877-L8881), so it's not checked. What was happening is that `embedding_renorm` was setting `torch.no_grad()` which only turns off the backwards mode AD so functorch's jvp tests were still using forward mode AD during the `embedding_renorm` call. This makes it so that we don't use forward mode during the embedding_renorm call Pull Request resolved: #78560 Approved by: https://github.com/soulitzer, https://github.com/albanD Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/ce7c7bb2a9b21c8a05eedc6b43b277654064a1c9 Reviewed By: b0noI Differential Revision: D36865377 fbshipit-source-id: acd33447f1213136019644bab6874ba38ff3836f

…orward mode AD (#78560)" Summary: This reverts commit ce7c7bb. Reverted #78560 on behalf of https://github.com/malfet due to broke XLA (on CI and trunk), see https://hud.pytorch.org/pytorch/pytorch/commit/ce7c7bb2a9b21c8a05eedc6b43b277654064a1c9 Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/d578197747db4efef7771bb2304d6bb54638eb73 Reviewed By: b0noI Differential Revision: D36882233 fbshipit-source-id: 1189a3da13824959f725cdcb784a01f3402faa78

samdow · 2022-06-03T19:13:39Z

@pytorchbot merge

…ode AD (#78560) (#78560) Summary: On functorch, we started seeing [embedding forward mode fail](pytorch/functorch#816). From looking at it, we figured out that recently [embedding got forward mode support enabled](369d9f4) and then doing forward mode with embedding and [max_norm doesn't work with gradcheck](https://github.com/pytorch/pytorch/blob/master/torch/testing/_internal/common_methods_invocations.py#L8877-L8881), so it's not checked. What was happening is that `embedding_renorm` was setting `torch.no_grad()` which only turns off the backwards mode AD so functorch's jvp tests were still using forward mode AD during the `embedding_renorm` call. This makes it so that we don't use forward mode during the embedding_renorm call Pull Request resolved: #78560 Approved by: https://github.com/soulitzer, https://github.com/albanD Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/b7cb4eae6b16ce312fb5e285aa82f227be6549a6 Reviewed By: b0noI Differential Revision: D36907413 Pulled By: samdow fbshipit-source-id: 52b6675ccb71287daed41f73a2d9f86b967d99b3

samdow requested a review from zou3519 May 31, 2022 19:16

samdow requested review from albanD and jbschlosser as code owners May 31, 2022 19:16

facebook-github-bot added the cla signed label May 31, 2022

zou3519 reviewed May 31, 2022

View reviewed changes

torch/nn/functional.py Outdated Show resolved Hide resolved

zou3519 requested review from soulitzer and zou3519 May 31, 2022 19:21

jbschlosser reviewed May 31, 2022

View reviewed changes

soulitzer approved these changes Jun 1, 2022

View reviewed changes

albanD approved these changes Jun 1, 2022

View reviewed changes

pytorchmergebot added the Merged label Jun 2, 2022

pytorchmergebot closed this in ce7c7bb Jun 2, 2022

samdow added release notes: nn release notes category topic: bug fixes topic category labels Jun 2, 2022

pytorchmergebot added the Reverted label Jun 2, 2022

samdow reopened this Jun 2, 2022

samdow force-pushed the fix_embedding branch from ee8030d to 23cc928 Compare June 2, 2022 19:37

make embedding work with forward mode AD

6f3457b

samdow added 2 commits June 3, 2022 12:38

use detach, add tests

cf56b29

fix xla failures

25309c2

samdow force-pushed the fix_embedding branch from 23cc928 to 25309c2 Compare June 3, 2022 16:38

pytorchmergebot closed this in b7cb4ea Jun 3, 2022

lkct mentioned this pull request Jun 17, 2023

Mis-annotated return for F._no_grad_embedding_renorm_ (also JIT related) #103800

Open

github-actions bot deleted the fix_embedding branch February 16, 2024 01:59

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fix embedding jvp support by making embedding_renorm ignore forward mode AD #78560

Fix embedding jvp support by making embedding_renorm ignore forward mode AD #78560

Uh oh!

samdow commented May 31, 2022

Uh oh!

facebook-github-bot commented May 31, 2022 •

edited

Loading

Uh oh!

Uh oh!

jbschlosser May 31, 2022

Uh oh!

soulitzer May 31, 2022

Uh oh!

samdow May 31, 2022

Uh oh!

samdow commented May 31, 2022

Uh oh!

soulitzer left a comment

Uh oh!

albanD left a comment

Uh oh!

samdow commented Jun 2, 2022

Uh oh!

github-actions bot commented Jun 2, 2022

Uh oh!

malfet commented Jun 2, 2022

Uh oh!

pytorch-bot bot commented Jun 2, 2022

Uh oh!

malfet commented Jun 2, 2022

Uh oh!

malfet commented Jun 2, 2022

Uh oh!

samdow commented Jun 3, 2022

Uh oh!

Uh oh!

Fix embedding jvp support by making embedding_renorm ignore forward mode AD #78560

Fix embedding jvp support by making embedding_renorm ignore forward mode AD #78560

Uh oh!

Conversation

samdow commented May 31, 2022

Uh oh!

facebook-github-bot commented May 31, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful links

✅ No Failures (0 Pending)

Uh oh!

Uh oh!

jbschlosser May 31, 2022

Choose a reason for hiding this comment

Uh oh!

soulitzer May 31, 2022

Choose a reason for hiding this comment

Uh oh!

samdow May 31, 2022

Choose a reason for hiding this comment

Uh oh!

samdow commented May 31, 2022

Uh oh!

soulitzer left a comment

Choose a reason for hiding this comment

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

samdow commented Jun 2, 2022

Uh oh!

github-actions bot commented Jun 2, 2022

Uh oh!

malfet commented Jun 2, 2022

Uh oh!

pytorch-bot bot commented Jun 2, 2022

Uh oh!

malfet commented Jun 2, 2022

Uh oh!

malfet commented Jun 2, 2022

Uh oh!

samdow commented Jun 3, 2022

Uh oh!

Uh oh!

facebook-github-bot commented May 31, 2022 •

edited

Loading