Skip to content

fix: preserve retained hooks during bridge helper cleanup - #1908

Open
emerardd wants to merge 2 commits into
TransformerLensOrg:devfrom
emerardd:fix/retained-hook-cleanup
Open

emerardd wants to merge 2 commits into
TransformerLensOrg:devfrom
emerardd:fix/retained-hook-cleanup

Conversation

@emerardd

Copy link
Copy Markdown
Contributor

Description

Preserve hooks explicitly retained with reset_hooks_end=False when subsequent Bridge helpers clean up their temporary hooks.

Retained hooks keep their lexical context level after a helper exits. A later sibling hooks(), run_with_hooks(), or run_with_cache() call reuses that level, so level-based cleanup can silently remove those retained hooks. In run_with_cache(), this can also remove retained hooks on unrelated hook points. This is an uncovered lifecycle case related to #1636 and #1638: directly registered hooks are already protected, but hooks intentionally retained from an earlier helper are not.

Each helper now records the exact LensHandle instances it registers and removes only its own temporary handles. Registration-failure cleanup and early-stop hooks use the same ownership rule. Public hook APIs, lexical context levels, explicit level-based resets, permanent-hook behavior, and context-clearing behavior remain unchanged.

Regression tests cover all three helpers, forward/backward hooks, same/unrelated hook points, nested retained scopes, exceptions, partial registration failures, caller-owned same-level hooks, and early-stop cleanup. A tiny Native model test also verifies that a retained intervention continues to affect logits after an unrelated cache call and that an explicit reset restores the baseline.

Type of change

  • Bug fix (non-breaking change which fixes an issue)

Validation

  • Before the fix, the retained-hook regressions produced 27 failed, 13 passed (24 other tests deselected).
  • After the fix, the focused orchestration, hook reset, alias resolution, caching, batch-dimension, stop-layer, and hook-gating tests produced 290 passed, 1 skipped, 3 xfailed. No skip or xfail was added.
  • uv run --no-sync mypy .: success, no issues in 388 source files.
  • Scoped pycln, isort, and Black checks passed, along with git diff --check. These are the Makefile-equivalent tools used on native Windows without GNU Make.

The complete unit-test suite and GPU tests were not run. The focused run emitted existing stub-adapter alias and dependency deprecation warnings; this patch adds no warning sites. No documentation update is needed because this restores the existing retained-hook contract without changing the public API.

Checklist

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@jlarson4 jlarson4 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks pretty solid! Just one note on an edge case that needs addressing when you have a moment. Thanks @emerardd

# `level` keeps this to the hooks added above — hooks the caller
# attached before the call survive.
self.remove_all_hook_fns(level=context_level)
for hp, direction, handle in added_hooks:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With stop_at_layer set and reset_hooks_end=False, the stop hook is kept along with the caller's hooks. Before this change, the next run_with_cache cleared it as a side effect. Now nothing does, so every later forward (including a plain bridge(tokens)) silently returns the residual stream at that layer instead of logits. Could the stop hook be removed unconditionally, since it's internal rather than something the caller asked to keep, with the reset_hooks_end=False case added to the stop-at-layer test?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for catching this — you're right that retaining requested hooks must not retain the internal stop hook. run_with_cache() now tracks its stop handle separately and always removes that exact handle in finally, including when reset_hooks_end=False or the forward raises. Requested cache hooks and previously retained caller hooks still survive according to the retention setting.

I expanded the stop-at-layer regression to cover both reset settings and both execution helpers. It checks that the stopped call returns the residual stream, the internal stop hook is gone, requested hooks remain, and both a subsequent plain forward and a subsequent cache call return the baseline logits. Additional exception-path tests cover a forward failing before the stop point.

Validation: 294 passed, 1 skipped, 3 xfailed across the focused hook orchestration/reset/alias/cache/stop-layer/gating tests; no new skips or xfails. Scoped pycln/isort/Black checks and full mypy (388 source files) also pass. The complete test suite was not rerun.

This branch has not been deployed

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants