fix(hooks): preserve original forward signature in register_hook#14237
Open
sergioperezcheco wants to merge 1 commit into
Open
fix(hooks): preserve original forward signature in register_hook#14237sergioperezcheco wants to merge 1 commit into
sergioperezcheco wants to merge 1 commit into
Conversation
HookRegistry.register_hook wraps the module forward in a (module, *args, **kwargs) closure and bound it via functools.update_wrapper pointing __wrapped__ at that closure. As a result inspect.signature(module.forward) resolved to (module, *args, **kwargs) instead of the real parameters, and torch.export/torch.compile failed with "missing a required argument: 'module'". Point __wrapped__ (via update_wrapper's wrapped argument) at the original forward captured before it is overwritten, so inspect.signature follows the chain back to the real signature. The chain works for multiple stacked hooks because each wrapper's __wrapped__ points at the previous (also-correct) wrapper. Call-time behaviour is unchanged. Regression test in tests/hooks/test_hooks.py asserts the signature and __name__ are preserved after registering one and two hooks. Fixes huggingface#14135
github-actions
Bot
added
fixes-issue
tests
hooks
size/S
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #14135.
HookRegistry.register_hook wraps the module forward through a (module, *args, **kwargs) closure and previously bound the result with functools.update_wrapper pointing wrapped at that closure. Because inspect.signature follows wrapped, the wrapped forward exposed (module, *args, **kwargs) instead of the real parameters, which broke torch.export/torch.compile with "missing a required argument: 'module'" and hid the real signature from tooling.
The fix points wrapped at the original forward captured before it is overwritten, so inspect.signature resolves back to the real parameters. This chains correctly when multiple hooks are stacked, since each wrapper's wrapped points at the previous (now also correctly-wrapped) wrapper. Runtime behaviour is unchanged.
Verified with the issue's reproducer: signature is preserved after one and two hooks, and torch.export.export succeeds where it previously failed. Added a regression test in tests/hooks/test_hooks.py asserting both the signature and name are preserved.
Assisted-by: Hermes Agent