Skip to content

Preserve HookRegistry forward signatures - #14143

Open
momomuchu wants to merge 1 commit into
huggingface:mainfrom
momomuchu:fix/hookregistry-forward-signature
Open

Preserve HookRegistry forward signatures#14143
momomuchu wants to merge 1 commit into
huggingface:mainfrom
momomuchu:fix/hookregistry-forward-signature

Conversation

@momomuchu

Copy link
Copy Markdown

Bug

HookRegistry.register_hook wrapped module.forward with a generic callable. After hook registration, inspect.signature(module.forward) reported:

(module, *args, **kwargs)

instead of the original model forward signature.

This broke torch.export keyword binding because kwargs such as hidden_states and timestep were bound against the wrapper signature.

Root cause

The replacement forward used functools.partial plus functools.update_wrapper, but it did not carry the original bound forward signature onto the installed wrapper.

Fix

Capture inspect.signature(forward) before wrapping, then assign it to self._module_ref.forward.__signature__ after installing the wrapped partial.

This keeps introspection aligned with the model API while leaving hook execution unchanged.

Tests

  1. Added a regression test that a single hook registration preserves inspect.signature(model.forward).
  2. Added a regression test that torch.export.export still accepts keyword inputs after hook registration.
  3. Added regression tests for repeated instrumentation: after registering two hooks, inspect.signature(model.forward) still equals the original, and torch.export.export still binds keyword inputs. This covers the case where a second register_hook could otherwise capture the wrapper signature instead of the model signature.
  4. Ran the relevant hooks suite:
DIFFUSERS_TEST_DEVICE=cpu python -m pytest -q tests/hooks/test_hooks.py
  1. Ran targeted ruff and format checks on the changed files.

Fixes #14135.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

HookRegistry.register_hook breaks module forward signature

1 participant