[Bug Report] PositionEmbeddingsAttentionBridge replaces eager_attention_forward in transformers' gemma-2 and gemma-3 modules for the rest of the process #1795

Human Preface (@speediedan)

We've contributed and thanked you before, but thanks again for TransformerLens! I'm sure you're aware how important it remains to the open-source AI interp community but it's integral to the AI world model analysis and tuning framework we're building (interpretune.org) so thanks again!

We've reviewed and human-validated the rest of this description below but it is largely claude-opus-5 written. Let us know if you want more human clarification/feedback/revision of it. Thanks!

Describe the bug

transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py::_setup_eager_attention_hook_wrapper
assigns a hook-firing wrapper over the module-level eager_attention_forward of
transformers.models.gemma2.modeling_gemma2 and transformers.models.gemma3.modeling_gemma3 the first time any
PositionEmbeddingsAttentionBridge is constructed. That component backs llama, qwen2/3, gemma 1/2/3, olmo, phi,
mixtral and about twenty other architectures, so building a bridge for a llama model patches both gemma modules.
The wrapper is installed once per process, is never removed (deleting the bridge does not restore it), and carries
no functools.wraps metadata, so nothing downstream can see through it.

Two consequences outside TransformerLens:

1. Any code that inspects the function bound at gemma's attention_interface call site sees TransformerLens's wrapper
   instead of transformers' function. nnsight's recursive .source tracing parses that function's source, and the
   wrapper's body has no nn.functional.dropout(...) call, so the operation node
   ...self_attn.source.attention_interface_0.source.nn_functional_dropout_0 that circuit-tracer uses as the attention
   pattern location for gemma models does not exist once any such bridge has been built. circuit-tracer's nnsight
   backend on gemma-3 then fails with AttributeError: 'SourceEnvoy' object has no attribute 'nn_functional_dropout_0'
   for the rest of the process, and only in processes where a bridge was built first, which makes it look like a
   flaky test in every downstream suite that mixes the two.
2. The wrapper runs for every gemma-2/3 forward in the process, bridged or not (it looks the bridge up by module id and
   passes through when none is registered), so plain HF gemma models pay the lookup on every attention call after any
   bridge has existed.

Code example

Measured with transformer-lens 3.5.1 (the patch is unchanged on main as of v3.8.1), transformers 5.16.1, CPU only:

import inspect
import transformers.models.gemma3.modeling_gemma3 as g3
import transformers.models.gemma2.modeling_gemma2 as g2
from transformer_lens.model_bridge import TransformerBridge

def state(tag):
    for name, mod in (("gemma3", g3), ("gemma2", g2)):
        f = mod.eager_attention_forward
        print(tag, name, f"{f.__module__}.{f.__qualname__}", "wrapped" if hasattr(f, "__wrapped__") else "no __wrapped__",
              "dropout calls in source:", inspect.getsource(f).count("dropout("))

state("before")
bridge = TransformerBridge.boot_transformers("google/gemma-3-270m", device="cpu")
state("after bridge")
del bridge
state("after del")

Output:

before gemma3 transformers.models.gemma3.modeling_gemma3.eager_attention_forward no __wrapped__ dropout calls in source: 1
before gemma2 transformers.models.gemma2.modeling_gemma2.eager_attention_forward no __wrapped__ dropout calls in source: 1
after bridge gemma3 transformer_lens.model_bridge.generalized_components.position_embeddings_attention._setup_eager_attention_hook_wrapper.<locals>.hooked_eager_attention_forward no __wrapped__ dropout calls in source: 0
after bridge gemma2 transformer_lens.model_bridge.generalized_components.position_embeddings_attention._setup_eager_attention_hook_wrapper.<locals>.hooked_eager_attention_forward no __wrapped__ dropout calls in source: 0
after del gemma3 transformer_lens.model_bridge.generalized_components.position_embeddings_attention._setup_eager_attention_hook_wrapper.<locals>.hooked_eager_attention_forward no __wrapped__ dropout calls in source: 0
after del gemma2 transformer_lens.model_bridge.generalized_components.position_embeddings_attention._setup_eager_attention_hook_wrapper.<locals>.hooked_eager_attention_forward no __wrapped__ dropout calls in source: 0

A llama bridge gives the same result for the two gemma modules (the patch is unconditional on the bridged
architecture); modeling_llama.eager_attention_forward itself is untouched.

Expected behavior

Building a bridge should not change what other libraries see when they inspect transformers' gemma modules, and
should leave no trace once the bridge is gone. Options, in rough order of how much they change:

• Register the hooked function as an attention implementation (AttentionInterface.register("tl_hooked_eager", ...))
  and select it on the bridged model's config, so the patch is per-model rather than per-process, and ordinary
  _attn_implementation="eager" models keep transformers' function.
• If the module-level assignment must stay, make it reversible (restore in the bridge's teardown / a context manager,
  reference-counted across live bridges), apply it only to the modeling module of the architecture being bridged,
  and decorate the wrapper with functools.wraps(original) so inspect.unwrap reaches the original.

System Info

• transformer-lens 3.5.1 (checked main at v3.8.1: position_embeddings_attention.py still assigns
  gemma2_module.eager_attention_forward / gemma3_module.eager_attention_forward)
• transformers 5.16.1, torch 2.x, Python 3.13, Linux, CPU reproduction

Additional context

Found through circuit-tracer's nnsight backend, whose attention-pattern location for gemma models is resolved
through nnsight's recursive source tracing; the failure appears only in processes where a bridge was built first.