Add hooks for register_buffer/module/parameter #86148

antoniojkim · 2022-10-03T19:23:49Z

As described in the issue, this PR adds hooks to be run when register_parameter, register_buffer and register_module are called.

Fixes #85837

cc @albanD @mruberry @jbschlosser @walterddr @kshitij12345 @saketh-are

pytorch-bot · 2022-10-03T19:23:51Z

🔗 Helpful Links

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

📄 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

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

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

albanD · 2022-10-03T19:36:47Z

Hi @antoniojkim thanks for sending a PR for this!

we review changes to nn.Module very carefully and in particular changing the direct setting of _modules to a register_module is something we should make sure is safe.
could you split the formatting change in a separate PR? That would make it a lot easier to review this one.

antoniojkim · 2022-10-03T19:42:04Z

* could you split the formatting change in a separate PR? That would make it a lot easier to review this one.

Would this PR pass the linting check in CI without the proper formatting?

albanD · 2022-10-03T19:46:45Z

Would this PR pass the linting check in CI without the proper formatting?

No. But we can merge the formatting PR quickly and so it will disappear from this PR.

antoniojkim · 2022-10-03T19:52:01Z

I've undone all the formatting changes

changing the direct setting of _modules to a register_module is something we should make sure is safe.

Turns out that the behaviour is slightly different. So, I've reverted it back to how it was before. There is now multiple copies of the code that calls the buffer registration and module registration hooks

antoniojkim · 2022-10-03T22:48:46Z

@pytorchbot rebase

pytorch-bot · 2022-10-03T22:48:48Z

You don't have permissions to rebase this PR, only people with write permissions may rebase PRs.

antoniojkim · 2022-10-03T22:49:05Z

@pytorchbot -h

pytorch-bot · 2022-10-03T22:49:07Z

PyTorchBot Help

usage: @pytorchbot [-h] {merge,revert,rebase,label} ...

In order to invoke the bot on your PR, include a line that starts with
@pytorchbot anywhere in a comment. That line will form the command; no
multi-line commands are allowed. 

Example:
    Some extra context, blah blah, wow this PR looks awesome

    @pytorchbot merge

optional arguments:
  -h, --help            Show this help message and exit.

command:
  {merge,revert,rebase,label}
    merge               Merge a PR
    revert              Revert a PR
    rebase              Rebase a PR
    label               Add label to a PR

Merge

usage: @pytorchbot merge [-g | -f MESSAGE | -l] [-r [{viable/strict,master}]]

Merge an accepted PR, subject to the rules in .github/merge_rules.json.
By default, this will wait for all required checks (lint, pull) to succeed before merging.

optional arguments:
  -g, --green           Merge when all status checks running on the PR pass. To add status checks, use labels like `ciflow/trunk`.
  -f MESSAGE, --force MESSAGE
                        Merge without checking anything. This requires a reason for auditting purpose, for example:
                        @pytorchbot merge -f 'Minor update to fix lint. Expecting all PR tests to pass'
  -l, --land-checks     Merge with land time checks. This will create a new branch with your changes rebased on viable/strict and run a majority of trunk tests _before_ landing to increase trunk reliability and decrease risk of revert. The tests added are: pull, Lint and trunk. Note that periodic is excluded. (EXPERIMENTAL)
  -r [{viable/strict,master}], --rebase [{viable/strict,master}]
                        Rebase the PR to re run checks before merging.  Accepts viable/strict or master as branch options and will default to viable/strict if not specified.

Revert

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

Revert a merged PR. This requires that you are a Meta employee.

Example:
  @pytorchbot revert -m="This is breaking tests on trunk. hud.pytorch.org/" -c=nosignal

optional arguments:
  -m MESSAGE, --message MESSAGE
                        The reason you are reverting, will be put in the commit message. Must be longer than 3 words.
  -c {nosignal,ignoredsignal,landrace,weird,ghfirst}, --classification {nosignal,ignoredsignal,landrace,weird,ghfirst}
                        A machine-friendly classification of the revert reason.

Rebase

usage: @pytorchbot rebase [-s | -b BRANCH]

Rebase a PR. Rebasing defaults to the stable viable/strict branch of pytorch.
You, along with any member of the pytorch organization, can rebase your PR.

optional arguments:
  -s, --stable          [DEPRECATED] Rebase onto viable/strict
  -b BRANCH, --branch BRANCH
                        Branch you would like to rebase to

Label

usage: @pytorchbot label labels [labels ...]

Adds label to a PR

positional arguments:
  labels  Labels to add to given Pull Request

linux-foundation-easycla · 2022-10-03T22:51:18Z

The committers listed above are authorized under a signed CLA.

✅ login: antoniojkim / name: Jae Hoon (Antonio) Kim (6ad0084d17dda3f25ec7247da6d168d5bd9d6c92)

albanD

Looks pretty good!
We would need a couple tests and update to the signature and that will be good to go!

albanD · 2022-10-04T14:30:32Z

torch/nn/modules/module.py

Could you name these with a "module" in the name to match the other hooks?

albanD · 2022-10-04T14:31:57Z

torch/nn/modules/module.py

Shouldn't the buffer be passed in?
Also for consistency, could you pass the module first: hook(module, name, buffer) -> None or new buffer

Same for the others below

yep, good catch. That's my bad for just copying the docstring over and not replacing all instances 😅

albanD

Sorry for the misunderstanding. I think we can do a small update for the signatures.
The calls will need to be updated as well to pass in the right arguments to the hooks.

albanD · 2022-10-05T15:00:38Z

torch/nn/modules/module.py

This should still be hook(module, name, buffer) -> None or new buffer

albanD · 2022-10-05T15:00:59Z

torch/nn/modules/module.py

I think this one should be hook(module, name, submodule) -> None or new submodule

albanD · 2022-10-05T15:01:19Z

torch/nn/modules/module.py

This one should be hook(module, name, param) -> None or new param

albanD

Signatures look good!
I think we only need two things now:

Update the call sites to pass in the right arguments
Add some basic tests for each of these hooks

albanD

SGTM!
Thanks for taking the time to add full testing!

albanD · 2022-10-12T17:59:33Z

@pytorchbot merge -g

pytorchmergebot · 2022-10-12T18:00:54Z

Merge started

Your change will be merged once all checks on your PR pass since you used the green (-g) flag (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

github-actions · 2022-10-12T21:11:25Z

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

@albanD

There is a bug in the implementation of the registration hooks introduced in #86148 whereby if the hook returns a tensor, then the short circuiting logic: ``` value = hook(self, name, value) or value ``` Raises an exception ``` RuntimeError: Boolean value of Tensor with more than one value is ambiguous ``` Fixing the logic so that it only checks to see if the value is `None` before overriding Fixes #85837 CC: @albanD @jbschlosser Pull Request resolved: #87369 Approved by: https://github.com/albanD

antoniojkim requested review from albanD and jbschlosser as code owners October 3, 2022 19:23

facebook-github-bot added the cla signed label Oct 3, 2022

pytorchbot added the open source label Oct 3, 2022

antoniojkim force-pushed the antoniojkim/new_registration_hooks branch from 2dd2179 to 546e93d Compare October 3, 2022 19:50

antoniojkim force-pushed the antoniojkim/new_registration_hooks branch from 546e93d to 6ad0084 Compare October 3, 2022 22:51

antoniojkim force-pushed the antoniojkim/new_registration_hooks branch from 6ad0084 to 127f0c9 Compare October 4, 2022 13:21

albanD reviewed Oct 4, 2022

View reviewed changes

albanD reviewed Oct 5, 2022

View reviewed changes

antoniojkim force-pushed the antoniojkim/new_registration_hooks branch from 5893b2d to 73d6ced Compare October 5, 2022 16:46

albanD reviewed Oct 5, 2022

View reviewed changes

bdhirsh added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Oct 6, 2022

antoniojkim added 2 commits October 11, 2022 06:30

Add hooks for register_buffer/module/parameter

d50db2b

Address PR comments

fb09e4e

antoniojkim force-pushed the antoniojkim/new_registration_hooks branch 3 times, most recently from 945dff2 to 86fa769 Compare October 11, 2022 22:26

Address PR comments and add tests

41607f6

antoniojkim force-pushed the antoniojkim/new_registration_hooks branch from 86fa769 to 41607f6 Compare October 12, 2022 14:49

albanD approved these changes Oct 12, 2022

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 12, 2022

pytorchmergebot added the Merged label Oct 12, 2022

pytorchmergebot closed this in 09a676f Oct 12, 2022

antoniojkim mentioned this pull request Oct 20, 2022

Fix registration hooks #87369

Closed

Add hooks for register_buffer/module/parameter #86148

Add hooks for register_buffer/module/parameter #86148

Uh oh!

Conversation

antoniojkim commented Oct 3, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Oct 3, 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/86148

✅ No Failures

Uh oh!

albanD commented Oct 3, 2022

Uh oh!

antoniojkim commented Oct 3, 2022

Uh oh!

albanD commented Oct 3, 2022

Uh oh!

antoniojkim commented Oct 3, 2022

Uh oh!

antoniojkim commented Oct 3, 2022

Uh oh!

pytorch-bot bot commented Oct 3, 2022

Uh oh!

antoniojkim commented Oct 3, 2022

Uh oh!

pytorch-bot bot commented Oct 3, 2022

PyTorchBot Help

Merge

Revert

Rebase

Label

Uh oh!

linux-foundation-easycla bot commented Oct 3, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

albanD Oct 4, 2022

Choose a reason for hiding this comment

Uh oh!

albanD Oct 4, 2022

Choose a reason for hiding this comment

Uh oh!

albanD Oct 4, 2022

Choose a reason for hiding this comment

Uh oh!

antoniojkim Oct 4, 2022

Choose a reason for hiding this comment

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

albanD Oct 5, 2022

Choose a reason for hiding this comment

Uh oh!

albanD Oct 5, 2022

Choose a reason for hiding this comment

Uh oh!

albanD Oct 5, 2022

Choose a reason for hiding this comment

Uh oh!

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

albanD commented Oct 12, 2022

Uh oh!

pytorchmergebot commented Oct 12, 2022

Merge started

Uh oh!

github-actions bot commented Oct 12, 2022

Uh oh!

Reviewers

Assignees

Labels

antoniojkim commented Oct 3, 2022 •

edited

Loading

pytorch-bot bot commented Oct 3, 2022 •

edited

Loading

linux-foundation-easycla bot commented Oct 3, 2022 •

edited

Loading