Add prepend argument to nn.Module hooks

[mrshenli]

Stack from ghstack (oldest at bottom):

• Add composable activation checkpointing #87664
• Add distributed composable API contract #87580
• -> Add prepend argument to nn.Module hooks #87370

cc @ezyang @gchanan

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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, 3 Pending

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

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: mrshenli (5b95241, 85ced4b, ceb933a, bb05bcd)

[albanD]

What is the exact use case for which you want this?
This is not very reliable in general as the next preprend hook you add will make your not the first one anymore.

cc @soulitzer related to hook ordering questions we were having.

[mrshenli]

What is the exact use case for which you want this?

We are working on converting some distributed features to hook-based solutions, so that we won't mess up FQNs as module wrapper did. To maintain the similar behavior as today DDP/FSDP, we plan to split their forward into top and bottom halves and install those as pre and post forward hooks. And the pre hook needs to be called before any existing hooks on the local module.

This is not very reliable in general as the next preprend hook you add will make your not the first one anymore.

This should be OK for now, as we expect DDP logic is added after local model preparations.

[albanD]

This should be OK for now, as we expect DDP logic is added after local model preparations.

Sure, but I'm not sure how to convey that to the end user of this flag in the doc :p (btw docs need updating).
Like we don't have a strict guarantee on the order of the backward hooks right now (which we want to figure out with @soulitzer btw). So I'm not sure what we're going to say here...

[soulitzer]

For hooks in general it feels like provide guarantees about ordering is something we'd want to do as long as it doesn't complicate the implementation, so the idea sounds OK.

For this case in particular though, communicating to the user about what "prepend" does can be tricky because hooks registered globally will still execute before those prepend, and I'm not sure we even document that anywhere currently?

[mrshenli]

Hey @albanD @soulitzer, thanks a lot for the comments!

For this case in particular though, communicating to the user about what "prepend" does can be tricky because hooks registered globally will still execute before those prepend,

Would I be correct that global hooks are mainly for debugging/profiling purposes. Or have we seen cases where program relies on global hooks to produce correct results?

and I'm not sure we even document that anywhere currently?

Any suggestions on which of the following might be better?

1. Add an underscore (i.e., _prepend) to make it private.
2. Add arg docstrings to make it official and fully explain the ordering of hooks

Edit: chatted with @albanD offline, since private args will still render in the doc, it's not really private. So let's avoid that.

[mrshenli]

[soulitzer]

Thanks for the updated documentation, looks great.

Add arg docstrings to make it official and fully explain the ordering of hooks

To explain the ordering of hooks fully, we may also want to think about the interaction of full_backward_pre_hooks with hooks registered to module outputs (but that might be out of the scope of this PR)

• FWIW the interaction seems to be: pre_hooks registered to the grad_fn of module outputs execute BEFORE prepended full_backward_pre_hooks, and post_hooks registered will execute AFTER. Not sure if it is problematic/relevant to your use case.

We might want to have some documentation regarding module hooks in general now, but that can be done in a follow up.

[github-actions]

This PR needs a label

If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

For more information, see https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[mrshenli]

@pytorchbot merge -g

[pytorchmergebot]

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