From 2e53914759dc246197c812e60db1ada7c5b5b375 Mon Sep 17 00:00:00 2001 From: emerard <113128214+emerardd@users.noreply.github.com> Date: Tue, 15 Sep 2026 23:11:11 +0800 Subject: [PATCH] fix: require eval mode for attribution patching --- tests/unit/tools/test_attribution_patching.py | 62 ++++++++++++++++++- .../tools/analysis/_model_state.py | 31 ++++++++++ .../tools/analysis/attribution_patching.py | 13 +++- .../tools/analysis/jacobian_lens.py | 29 +-------- 4 files changed, 105 insertions(+), 30 deletions(-) create mode 100644 transformer_lens/tools/analysis/_model_state.py diff --git a/tests/unit/tools/test_attribution_patching.py b/tests/unit/tools/test_attribution_patching.py index fe7339146..eda937a39 100644 --- a/tests/unit/tools/test_attribution_patching.py +++ b/tests/unit/tools/test_attribution_patching.py @@ -78,6 +78,7 @@ def __init__(self, *, dtype: torch.dtype = torch.float32) -> None: ) self.compatibility_mode = False self._weights_processed = False + self.forward_calls = 0 self.embed = nn.Embedding(D_VOCAB, D_MODEL, dtype=dtype) nn.init.normal_(self.embed.weight, std=0.2) self.hook_embed = HookPoint() @@ -88,6 +89,7 @@ def __init__(self, *, dtype: torch.dtype = torch.float32) -> None: self.ln_final = nn.Identity() self.unembed = nn.Linear(D_MODEL, D_VOCAB, bias=False, dtype=dtype) nn.init.normal_(self.unembed.weight, std=0.2) + self.eval() @property def hook_dict(self) -> dict[str, HookPoint]: @@ -118,6 +120,7 @@ def to_tokens(self, prompt: str) -> torch.Tensor: def forward( self, tokens: torch.Tensor, return_type: str | None = "logits" ) -> torch.Tensor | None: + self.forward_calls += 1 residual = self.hook_embed(self.embed(tokens)) for block in self.blocks: residual = block(residual) @@ -364,6 +367,7 @@ def __init__(self, d_model: int, n_heads: int, d_head: int, layer: int, dtype: t self.w_z = nn.Linear(d_model, n_heads * d_head, bias=False, dtype=dtype) self.w_o = nn.Linear(n_heads * d_head, d_model, bias=False, dtype=dtype) self.w_mlp = nn.Linear(d_model, d_model, bias=False, dtype=dtype) + self.dropout = nn.Dropout(p=0.5) for linear in (self.w_z, self.w_o, self.w_mlp): nn.init.normal_(linear.weight, std=0.2) self.hook_z = HookPoint() @@ -374,7 +378,8 @@ def __init__(self, d_model: int, n_heads: int, d_head: int, layer: int, dtype: t def forward(self, residual: torch.Tensor) -> torch.Tensor: batch, seq, _ = residual.shape z = self.hook_z(self.w_z(residual).reshape(batch, seq, self.n_heads, self.d_head)) - residual = residual + self.w_o(z.reshape(batch, seq, self.n_heads * self.d_head)) + attn_out = self.w_o(z.reshape(batch, seq, self.n_heads * self.d_head)) + residual = residual + self.dropout(attn_out) mlp_out = self.hook_mlp_out(self.w_mlp(residual)) return residual + mlp_out @@ -403,6 +408,7 @@ def __init__(self, *, dtype: torch.dtype = torch.float32) -> None: ) self.compatibility_mode = False self._weights_processed = False + self.forward_calls = 0 self.embed = nn.Embedding(D_VOCAB, D_MODEL, dtype=dtype) nn.init.normal_(self.embed.weight, std=0.2) self.hook_embed = HookPoint() @@ -413,6 +419,7 @@ def __init__(self, *, dtype: torch.dtype = torch.float32) -> None: self.ln_final = nn.Identity() self.unembed = nn.Linear(D_MODEL, D_VOCAB, bias=False, dtype=dtype) nn.init.normal_(self.unembed.weight, std=0.2) + self.eval() @property def hook_dict(self) -> dict[str, HookPoint]: @@ -425,6 +432,7 @@ def hook_dict(self) -> dict[str, HookPoint]: def forward( self, tokens: torch.Tensor, return_type: str | None = "logits" ) -> torch.Tensor | None: + self.forward_calls += 1 residual = self.hook_embed(self.embed(tokens)) for block in self.blocks: residual = block(residual) @@ -513,6 +521,57 @@ def test_attribution_patch_raises_on_batch_size_mismatch() -> None: attribution_patch(model, clean, corrupt, _metric_fn(answer=1, wrong=2)) +def test_attribution_patch_rejects_model_in_training_mode() -> None: + model = _NodeGraphToyBridge() + model.train() + tokens = torch.tensor([[1, 2, 3]]) + + with pytest.raises(ValueError, match=r"model\.eval\(\)"): + attribution_patch(model, tokens, tokens, _metric_fn(answer=1, wrong=2)) + assert model.training is True + assert model.forward_calls == 0 + + +def test_attribution_patch_rejects_nested_submodule_in_training_mode() -> None: + model = _NodeGraphToyBridge() + model.blocks[1].dropout.train() + tokens = torch.tensor([[1, 2, 3]]) + assert model.training is False + assert model.blocks[1].training is False + + with pytest.raises(ValueError, match=r"model\.eval\(\)"): + attribution_patch(model, tokens, tokens, _metric_fn(answer=1, wrong=2)) + assert model.training is False + assert model.blocks[1].training is False + assert model.blocks[1].dropout.training is True + assert model.forward_calls == 0 + + +def test_attribution_patch_rejects_hidden_original_model_training_mode() -> None: + model = _NodeGraphToyBridge() + original_model = nn.Sequential(nn.Linear(D_MODEL, D_MODEL)) + original_model.eval() + original_model[0].train() + model.__dict__["original_model"] = original_model + tokens = torch.tensor([[1, 2, 3]]) + + with pytest.raises(ValueError, match="original_model"): + attribution_patch(model, tokens, tokens, _metric_fn(answer=1, wrong=2)) + assert original_model.training is False + assert original_model[0].training is True + assert model.forward_calls == 0 + + +def test_attribution_patch_identical_inputs_are_exactly_zero_in_eval_mode() -> None: + model = _NodeGraphToyBridge() + tokens = torch.tensor([[1, 2, 3]]) + + result = attribution_patch(model, tokens, tokens, _metric_fn(answer=1, wrong=2)) + + assert result.node_scores + assert all(score == 0.0 for score in result.node_scores.values()) + + # --------------------------------------------------------------------------- # Commit 5 — linear-model reconstruction identity # --------------------------------------------------------------------------- @@ -670,6 +729,7 @@ def __init__(self, *, dtype: torch.dtype = torch.float32) -> None: for layer in range(N_LAYERS) ] ) + self.eval() def test_nonlinear_node_scores_read_the_corrupt_run_gradient() -> None: diff --git a/transformer_lens/tools/analysis/_model_state.py b/transformer_lens/tools/analysis/_model_state.py new file mode 100644 index 000000000..4ef787e75 --- /dev/null +++ b/transformer_lens/tools/analysis/_model_state.py @@ -0,0 +1,31 @@ +"""Shared model-state validation for analysis tools.""" + +from typing import Any + +import torch + + +def require_eval_mode(model: Any, *, operation: str) -> None: + """Reject training state anywhere in a wrapped model without mutating it.""" + training_modules: dict[int, str] = {} + roots = (("", model), ("original_model", getattr(model, "original_model", None))) + for prefix, root in roots: + if not isinstance(root, torch.nn.Module): + continue + for name, module in root.named_modules(): + if not module.training: + continue + qualified_name = ".".join(part for part in (prefix, name) if part) + training_modules.setdefault(id(module), qualified_name or "") + if not training_modules: + return + + names = list(training_modules.values()) + preview = ", ".join(names[:3]) + if len(names) > 3: + preview += f", and {len(names) - 3} more" + raise ValueError( + f"{operation} requires the model and all submodules to be in evaluation " + f"mode; found training mode at {preview}. Call model.eval() before running " + "the analysis." + ) diff --git a/transformer_lens/tools/analysis/attribution_patching.py b/transformer_lens/tools/analysis/attribution_patching.py index 1d5f17461..a0041ca4e 100644 --- a/transformer_lens/tools/analysis/attribution_patching.py +++ b/transformer_lens/tools/analysis/attribution_patching.py @@ -36,6 +36,8 @@ import torch +from transformer_lens.tools.analysis._model_state import require_eval_mode + MetricFn = Callable[[torch.Tensor], torch.Tensor] NamesFilter = Union[str, Sequence[str], Callable[[str], bool], None] @@ -460,6 +462,10 @@ def attribution_patch( reconstruction identity holds) and per-node scores are averaged across the batch before ranking. + The model and every submodule must be in evaluation mode. Separate clean and + corrupt forwards cannot produce meaningful activation differences if stochastic + training layers such as dropout remain active. + Args: model: A ``TransformerBridge`` (or compatible) exposing ``cfg.n_layers``, ``hook_dict``, and ``hooks()``. @@ -476,8 +482,9 @@ def attribution_patch( Raises: ValueError: if ``clean``/``corrupt`` are not 2D, hold a different number of - pairs, or a pair tokenizes to different lengths (activations must align - position-by-position). + pairs, a pair tokenizes to different lengths (activations must align + position-by-position), or the model or one of its submodules is in + training mode. """ del config # node granularity + ig_steps=1 only; enforced at construction. @@ -498,6 +505,8 @@ def attribution_patch( "Attribution patching aligns activations position-by-position." ) + require_eval_mode(model, operation="attribution_patch()") + node_hook_names = _required_hook_names(int(model.cfg.n_layers)) batch = int(clean.shape[0]) totals: dict[Node, float] = {} diff --git a/transformer_lens/tools/analysis/jacobian_lens.py b/transformer_lens/tools/analysis/jacobian_lens.py index a5b2d7102..34e7b5969 100644 --- a/transformer_lens/tools/analysis/jacobian_lens.py +++ b/transformer_lens/tools/analysis/jacobian_lens.py @@ -81,6 +81,7 @@ from tqdm.auto import tqdm from transformer_lens.ActivationCache import ActivationCache +from transformer_lens.tools.analysis._model_state import require_eval_mode from transformer_lens.tools.analysis.jacobian_lens_coordinate_patch import ( CoordinatePatch, solve_coordinate_patch, @@ -1701,7 +1702,7 @@ def fit( or layer indices, or if no prompt was long enough to fit on. """ _require_raw_bridge(model) - _require_eval_mode_for_fit(model) + require_eval_mode(model, operation="JacobianLens.fit()") if not isinstance(corpus, str) or not corpus.strip(): raise ValueError("corpus must be a non-empty provenance identifier") n_layers = model.cfg.n_layers @@ -1869,32 +1870,6 @@ def _require_raw_bridge(model: Any) -> None: ) -def _require_eval_mode_for_fit(model: Any) -> None: - """Reject stochastic training state without mutating the caller's model.""" - training_modules: Dict[int, str] = {} - roots = (("", model), ("original_model", getattr(model, "original_model", None))) - for prefix, root in roots: - if not isinstance(root, torch.nn.Module): - continue - for name, module in root.named_modules(): - if not module.training: - continue - qualified_name = ".".join(part for part in (prefix, name) if part) - training_modules.setdefault(id(module), qualified_name or "") - if not training_modules: - return - - names = list(training_modules.values()) - preview = ", ".join(names[:3]) - if len(names) > 3: - preview += f", and {len(names) - 3} more" - raise ValueError( - "JacobianLens.fit() requires the model and all submodules to be in " - f"evaluation mode; found training mode at {preview}. Call model.eval() " - "before fitting." - ) - - def _validate_metadata(metadata: Dict[str, Any]) -> None: """Reject values that ``torch.load(weights_only=True)`` cannot reload."""