Conversation
* Clean up warnings in CI * format fixes * regenerated lock file
* siglip vision tower issues * Removed expected failures
- Compare CPU fp32 Bridge logits against an independent eager Hugging Face reference. - Validate Bridge-native cached embeddings, block inputs, queries, normalization, and decoder outputs. - Cover padded batches and attention-mask behavior with fixed token IDs.
* docs: add Jacobian Lens decomposition demo * docs(jacobian_lens): synchronize demo captions with saved outputs * fix(demos): read RUN_FULL_ANALYSIS from environment variable Notebook hard-coded RUN_FULL_ANALYSIS = True, so the documented quick tier was only reachable by hand-editing the notebook. Read it from the RUN_FULL_ANALYSIS env var (default "true") so CI or callers can select the quick tier without touching the file.
…imitive) (#1741) * fix(jacobian_lens): support sparse decomposition on MPS * feat(jacobian_lens): add anchored coordinate-patching core - Model-free primitive for anchored edits over sparse J-space decompositions (PR1 commit 1 of 3) - solve_coordinate_patch edits one sparse coordinate (substitute overwrites the target, swap exchanges), appends an absent target at zero, blends by alpha (alpha=0 is a bit-exact no-op), and reconstructs over the anchored residual x - reconstruction. - CoordinatePatch is a frozen, fully detached report holding only frames, vectors, and diagnostics. - Reviewer constraints: Q2 overwritten_target_coordinate reports the exact discarded value and is None for swap or an absent target. - Q3/Q3-ext emit real non-fatal UserWarnings naming the basis condition number and the source/target cosine; the patch performs no inverse, so it warns and never raises even at the parallel extreme. - Q4 accepts a precomputed JSpaceDecomposition only after strict compatibility validation, yielding a patch identical to a fresh solve. - The nonedited-coordinate and residual max-delta witnesses are computed and enforced. - The shared _SWAP_WARN_COSINE threshold now lives in the decomposition leaf module; the swap_hooks rewire and intervention-module removal land in commit 2. * feat(jacobian_lens): expose coordinate patching on the Bridge - JacobianLens.coordinate_patch resolves a raw activation or a prompt+position exactly like decompose, maps the source/target tokens to dictionary-row indices, and calls the model-free solve_coordinate_patch. Offline only in PR1 — it returns a CoordinatePatch and installs no forward hooks - Shared-policy refactor (plan section 6): the near-parallel warn/raise helper and the _SWAP_ERROR_COSINE threshold now live beside _SWAP_WARN_COSINE in the decomposition leaf module, the common import both swap_hooks and the patch core already depend on. swap_hooks reads that one definition; its warn-at-0.99 / raise-at-0.999 behavior and messages stay byte-identical. - The temporary jacobian_lens_intervention module is removed - Tests: raw-activation and prompt/position paths agree; token inputs resolve to the correct dictionary-row indices. - A planted two-atom positive swap matches the pseudoinverse oracle (the pseudoinverse stays out of the core). - An unfitted layer raises; the core's near-parallel and conditioning warnings propagate unchanged through the wrapper; existing swap_hooks warn/raise coverage is unchanged after the constant move. * feat(jacobian_lens): export coordinate patching with docs and integration test Complete PR1 for J-space coordinate patching: - Export CoordinatePatch and solve_coordinate_patch from transformer_lens.tools.analysis and note coordinate patching in the package docstring tool list. - Add a cached-GPT-2 integration test asserting algebraic invariants only (finite output, residual preservation, an exact alpha=0 no-op, and unchanged non-selected coordinates) rather than a token flip. - Document the coordinate_patch API: substitute/swap semantics, target append, alpha interpolation, the offline/vocabulary-scale scope, and the warning policy (both the conditioning and near-parallel warnings are non-fatal and name the measured condition number / cosine, and unlike swap_hooks the patch core only warns at the parallel extreme). * fix(jacobian_lens): correct edit delta and make postcondition witnesses independent - Compute delta as basis @ (coordinates_after - coordinates_before) instead of basis @ coordinates_after - reconstruction_before, removing the ~2e-6 alpha-independent recompute floor and the discontinuity at alpha=0. - Retire the alpha=0 special case: the new delta is exactly zero and strictly proportional to alpha, so patched is continuous through the origin. - Anchor reconstruction_after to reconstruction_before + delta so the residual x - reconstruction_before is preserved exactly. - Recompute the residual witness from the dictionary rows (patched - dictionary[support_after].T @ coordinates_after vs residual) so it catches a basis/coordinate misalignment ins catches a basis/coordinate misalignment ins catches a basis/coordinate misalignment ins es with alpha (no floor). * refactor(analysis): extract shared MPS→CPU device-fallback helper - Add `_linalg_on_cpu_if_mps` in jacobian_lens_decomposition.py, centralizing the detect-device → `.cpu()` → op → restore idiom for torch.linalg routines MPS lacks kernels for (pinv, svdvals). - Replace all three duplicated call sites: `_validate_decomposition` (pinv projection) in jacobian_lens_coordinate_patch.py, the j_space_component pinv block in get_sparse_decomposition (jacobian_lens_decomposition.py), and `_singular_values` in projection_kernel.py. - Home module imports nothing intra-analysis, so the two consumers import it without introducing an import cycle. - Pure refactor, no behavior change: CPU path is bit-identical and the MPS branch is preserved. 387 unit tests across the touched tools pass; mypy clean. * fix(jacobian_lens): validate a supplied decomposition against the activation - Add an x-coupled check in _validate_decomposition: reuse the solver's NNLS KKT stationarity test so a self-consistently scaled (reconstruction, coordinates) pair that does not fit the activation is rejected instead of silently anchoring the edit to a wrong residual. - Tolerances mirror the solver's result-dtype certification (sqrt(eps) of the float32 coordinates) so genuine float32-rounded decompositions still pass. - Cover the case with test_rejects_scaled_but_self_consistent_decomposition. * fix(jacobian_lens): scale numerical tolerances at model scale - C2: in _basis_diagnostics, scale rank_rtol by the edit basis column count instead of max(shape) (the d_model row count), so a genuinely full-rank basis keeps a finite condition number past d_model >= ~2896 rather than reporting rank deficiency (condition number=inf) on Llama/Mistral/Gemma-scale bases. - C8: give _close an optional condition_number that scales the fixed 32-eps tolerance. _validate_decomposition measures the selected span's condition once (new _condition_number helper, MPS-safe) and relaxes the compatibility checks by it, so an exactly-valid decomposition whose activation lies along the span's near-null direction is no longer rejected as incompatible. - Add regression tests for a full-rank basis at large d_model and a high-condition decomposition that now validates. * feat(coordinate-patch): sharpen edit-result semantics and reporting - Add `target_was_selected` field distinguishing "pursuit never considered the target" from "pursuit selected it but assigned a zero coordinate" — the distinction `target_was_appended` alone couldn't express (C6) - Warn and add `coordinates_after_nonnegative` field when `alpha` outside [0, 1] extrapolates a coordinate negative, out of the c >= 0 pursuit frame (C5) - Document `overwritten_target_coordinate` as alpha-independent by design: it reports the coordinate a substitute targets for replacement, not the blended outcome, so it stays non-None even at alpha=0 where nothing is actually discarded (C10) * fix(jacobian-lens): reject boolean token inputs in the wrapper - `_to_token_ids` screened non-string tokens through bare `int(token)`, so `True`/`False` silently coerced to id `1`/`0` and later failed downstream with a misleading "source_idx=1 is not in the active support" error. - Raise `ValueError` on a bool token before the `int()` call, at the wrapper's token-resolution boundary (core `_validate_index` already rejected bools, but only after resolution). - Add `test_to_token_ids_rejects_bool` (True and False) asserting the clear error. * docs(jacobian-lens): give numeric values for conditioning thresholds - State basis_condition_number warning threshold: 1 / sqrt(float32 eps) (≈ 2896) - State _SWAP_WARN_COSINE (0.99) and _SWAP_ERROR_COSINE (0.999) values inline - Keep existing field pointers (source_target_cosine, basis_condition_number)
* Fix Cohere compatibility logit scaling * Handle missing Cohere runtime logit scale --------- Co-authored-by: jlarson4 <jonahalarson@comcast.net>
…tors (#1723) * feat(backward_lens): add gradient factorization contracts - Reconstruct linear weight gradients from forward inputs and output VJPs. - Support GPT-2 Conv1D and torch.nn.Linear weight layouts. - Add detached CPU contracts, validation, error metrics, and signed vocabulary rankings. - Cover numerical, validation, overflow, ownership, and ranking edge cases. * feat(backward_lens): capture GPT-2 MLP gradient factors - Validate raw GPT-2 Bridge models, prompts, targets, layers, and original Conv1D weights. - Capture both MLP projections with one forward pass and one autograd call. - Reconstruct exact weight gradients and return detached CPU-owned factors. - Preserve model state, RNG, gradients, and existing hooks with comprehensive integration coverage. * feat(backward_lens): project GPT-2 MLP gradient factors into vocabulary space - Add public BackwardLens.analyze(prompt, target_token, layers, normalized=False) returning detached, CPU-owned result dataclasses. - Project residual-width factors through fresh ln_final and unembed, with an optional Normalized Logit Lens for low-norm factors. - Expose factor norms, zero-norm masks, signed top/bottom vocabulary rankings, token decoding, and raw-gradient target ranks. - Export the public result contracts and cover the API with model-free unit tests and GPT-2 integration tests. * docs(backward_lens): document GPT-2 gradient factor analysis - Derive the FF1 and FF2 gradient factorizations, vocabulary projections, and tensor-shape contracts. - Document the public API, raw Bridge and single-token restrictions, state-safety guarantees, and error behavior. - Explain raw versus normalized projections, zero-factor handling, gradient-versus-SGD signs, interpretation limits, and troubleshooting. - Add the guide to the documentation index and cite Katz et al. without copying external code or assets. * docs(backward_lens): add executed GPT-2 gradient factor demo - Add an executed GPT-2-small walkthrough of MLP gradient factorization and vocabulary projection. - Visualize reconstruction error, numerical rank, layer-position token directions, VJP norms, and target ranks. - Compare raw and normalized projections, expose gradient-update sign semantics, mark near-zero signals, and sort cumulative reconstruction by contribution. - Save reproducible outputs and register all ten notebook cells in CI and the Makefile notebook test target. * fix(backward_lens): stabilize demo notebook validation - Suppress the narrow Typeguard instrumentation warning emitted during fresh TransformerLens imports. - Display target ranks as stable vocabulary percentiles while retaining exact ranks in the analysis data. - Regenerate all notebook outputs and verify all ten nbval cells pass. * fix(backward_lens): stabilize rank percentile output * fix(backward_lens): support beartype test instrumentation * fix(backward_lens): support beartype coverage instrumentation * fix(backward_lens): ignore setup-cell notebook output * feat(backward_lens): bound vocabulary readout retention * fix(backward_lens): harden capture preconditions * refactor(backward_lens): align tensor validation contracts * test(backward_lens): strengthen gradient behavior coverage * docs(backward_lens): clarify factor and token semantics * fixing export ordering * fix(backward_lens): restore weight via in-place copy instead of functional_call - swap torch.func.functional_call reparametrization for weight.copy_(updated_weight) under no_grad - restore original weight via try/finally so cleanup happens even if assertion fails - functional_call left layer 11 c_proj weight overridden in the bridge's original_model, breaking the three public_* fixtures later in this module (coverage-tier failure; Compatibility jobs don't run this tier) --------- Co-authored-by: jlarson4 <jonahalarson@comcast.net>
generate_stream(return_type="str") decoded tokens[0] only, so a batched stream yielded the first sequence's text and the rest of the batch was unreachable — generation itself was fine, only the decode step dropped the other rows. Decode each row and keep generate()'s unwrap convention: a bare string for a one-row batch, one string per row otherwise. Covered at the unit tier with a stubbed token stream and at the integration tier against distilgpt2. Fixes #1756
* fix: preserve batched hf generation context * fix: respect model-specific generation padding
* feat(jacobian_lens_coordinate_patch): add per-position patch loop core - Add solve_coordinate_patch_positions: model-free loop applying solve_coordinate_patch independently to every (batch, position) pair in a [batch, num_positions, d_model] chunk - Support an optional caller-owned decomposition_cache keyed (layer, batch_idx, position); a hit skips get_sparse_decomposition, a miss solves once and stores - Fail fast (no try/except) when a source is inactive at any pair, so no partial write reaches the activation tensor - Validate 3-D activations and matching position_labels length - Add unit tests for offline parity, batch independence, fail-fast, cache miss/hit behavior, and label-length validation * test(tools): extract shared J-lens toy-bridge fixtures into conftest - move D_MODEL/N_LAYERS/D_VOCAB/SEQ_LEN/SKIP_FIRST/CORPUS, _ToyBlock, _CausalSumBlock, _ToyTokenizer, _ToyBridge, _NotABridge, _lens, and the toy_model fixture from test_jacobian_lens.py into conftest.py - import the shared symbols back into test_jacobian_lens.py and remove the now-unused contextmanager and HookPoint imports - centralize the shared test setup so additional hook tests can reuse the same fixtures without duplicating roughly 90 lines of test code * feat(jacobian_lens): expose coordinate_patch_hooks on the Bridge - Add JacobianLens.coordinate_patch_hooks, a forward-hook variant of the offline coordinate_patch primitive, following the swap_hooks builder pattern - Solve one J-space coordinate patch per (batch, position) pair at each layer via solve_coordinate_patch_positions, with an optional caller-owned decomposition_cache keyed (layer, batch_idx, position) - Require positions explicitly and reject identical source/target tokens; fail fast on any inactive source rather than partially patching a batch - Warn once per call naming the layer x position count that performs a live vocabulary-scale solve on every cache miss - Add a dedicated test file covering shape parity, warning-once, cache hit/miss, oracle parity with offline coordinate_patch, and uncaught error and warning propagation through the hook * feat(jacobian_lens): export and document coordinate_patch_hooks Export solve_coordinate_patch_positions from tools.analysis and document the distinction between offline and dynamic/hooked coordinate patching. Replace the stale dynamic-patching claim in jacobian_lens_fitting.md and add documentation covering required positions, caller-owned decomposition_cache, per-pair fail-fast behavior, and the once-per-call cost warning. Add a cached GPT-2 integration test verifying alpha=0 is an exact no-op and untouched positions remain bit-identical. * fix(jacobian_lens): drop leaked sparse_probing exports from analysis __init__ Commit 02ffc27 accidentally added sparse_probing imports and __all__ entries to transformer_lens/tools/analysis/__init__.py without committing the sparse_probing module itself. On CI (which only checks out tracked files) importing transformer_lens.tools.analysis raised ModuleNotFoundError, breaking package import and failing every job that imports transformer_lens (unit, docstring, compatibility, benchmark, coverage, notebooks). Remove the out-of-scope sparse_probing exports; that work belongs to its own PR. * chore(jacobian_lens): drop unrelated sparse-probing and cursor changes from PR scope - Restore .cursor/rules/transformerlens.mdc deleted on this branch (unrelated to coordinate-patch hooks) - Remove sparse_probing import + test from tests/mps/test_mps_basic.py that broke tests/mps collection with ModuleNotFoundError (accidentally picked up from another branch) - Remove content/sparse_probing toctree line from docs/source/index.md referencing a nonexistent page * fix(jacobian_lens): straight-through gradient in coordinate-patch position loop - Scatter `x + patch.delta` instead of the detached `patch.patched` in solve_coordinate_patch_positions, keeping x's graph; bitwise identical forward for float32 activations. - Add test_solve_coordinate_patch_positions_preserves_gradient_to_patched_position: asserts finite nonzero grad at the patched position (exactly 0 before) and bitwise-equal forward output. * test(jacobian_lens): pin batch component of decomposition_cache key - Add test_solve_coordinate_patch_positions_populates_batch_component_of_cache_key: two distinct batch rows share one cache at the same (layer, position); assert both (layer, 0, position) and (layer, 1, position) populate with different active supports. - Closes the gap where the batch-agnostic [1, 1, d_model] fixtures let a collapsed (layer, 0, position) key pass the whole unit tier; the new test fails when the key drops batch_idx. * docs+test(jacobian_lens): document top-k support precondition and cover plural-layer install - Strengthen coordinate_patch_hooks Raises note: source must be in every patched position's top-k active support after earlier band hooks edit the residual, not merely active on a clean pass; substitute/swap remove the source coordinate, so stacking layers/positions makes it progressively harder. - Add the first plural (two-layer) end-to-end hook test, pinning the per-closure layer=layer/dictionary=dictionary binding via distinct per-layer dictionaries and per-layer cache keys. * docs(jacobian_lens): correct dynamic-hook generate() claim and cache-reuse caveat - Drop the model.generate(...) claim from the dynamic coordinate-patch hooks section: on GPT-2 with the published lens, generate re-decomposes each freshly generated token whose top-k support generally lacks the source concept, so fail-fast aborts generation. Keep run_with_hooks(...). - Document that the (layer, batch_idx, position) cache key uses the chunk-local position index, so decomposition_cache is valid only across identical-chunking passes; warn against reuse across decode steps (use_past_kv_cache prefill vs per-step shapes collide) or differently-shaped prompts.
…ansformerBridge (#1750) * feat(attribution_patching): names-filtered gradient-cache substrate Add the substrate for attribution patching on TransformerBridge: a single forward pass plus a manual metric backward that caches activations together with the gradient of a custom (non-scalar) metric with respect to each of them. run_with_cache(incl_bwd=True) only backpropagates the model's own scalar output, so the metric backward is done by hand. Gradients are retained only for hook points passing names_filter, keeping gradient-cache memory bounded to the hook families an analysis actually reads. The module docstring pins the denoising sign/direction convention (gradient from the corrupt run, estimate points toward the clean activation) that the node/edge scoring in follow-on commits builds on. Tested on a tiny fully-linear TransformerBridge subclass: the gradient cache covers exactly the filtered hook names and matches the closed-form linear gradient. * feat(attribution_patching): typed computational-graph node model - Add frozen Node dataclass keyed by (kind, layer, position, head) with per-kind field invariants enforced in __post_init__ (embed / attn-head-out / mlp-out) - Map each node to its standard TransformerBridge hook alias via Node.hook_name (hook_embed, blocks.{l}.attn.hook_z, blocks.{l}.hook_mlp_out) - Add enumerate_nodes to build the full explicit node graph from the hook graph, inferring seq_len and n_heads from cache shapes and n_layers from cfg - Raise (not skip) when a required hook point is missing from the cache, per Risk 1 explicit-graph guard - Add unit tests for node key validation, enumeration keys, and the missing-hook raise * feat(attribution_patching): config + result API - Add EdgeAttributionConfig with orthogonal granularity + ig_steps axes; remove the overlapping method axis to keep the config invalid-state-free - Default to the only executable path (granularity=node, ig_steps=1); document that the proposal's EAP-IG ig_steps=5 default lands in PR3 - Guard unimplemented paths: granularity=edge raises NotImplementedError pointing to PR2, while ig_steps>1 points to PR3 - Add AttributionResult with node_scores and magnitude-ranked top_nodes(k); declare edge_scores and top_edges() for API stability, with top_edges() raising until PR2 - Rely on beartype's Literal enforcement for granularity instead of a redundant runtime check - Add unit tests for config defaults, ig_steps validation, both NotImplementedError guards, top_nodes ranking order, and the top_edges guard * feat(attribution_patching): node attribution patch entry point - Add attribution_patch(model, clean, corrupt, metric_fn, config): two forwards + one manual backward per pair, computing node effects from clean-corrupt activations and the metric gradient - Validate token-length and batch-size parity between clean and corrupt inputs; raise on mismatch - Score each pair independently and average per-node scores across the batch - Pin the denoising convention (gradient from corrupt run, patch toward clean) in the docstring, with PR5 mapping it to the pinned oracle - Add _node_effects helper to contract the feature dimension per node position/head - Add _NodeGraphToyBridge fixture with attn.hook_z + hook_mlp_out, covering finite/shaped scores, batch averaging, and both parity guards * test(attribution_patching): linear-model reconstruction identity - Add exact-equality reconstruction test on a complete cut: summed embed-node scores reconstruct m(clean) - m(corrupt) for a linear model - Document why summing all nodes double-counts, since each node's full gradient re-counts paths through its upstream writes; the embed layer is the input-side complete cut - Add single-node exact-patch tests across embed, attn-head, and mlp families, verifying that corrupt->clean patching changes the metric by exactly the node score - Pin the denoising sign convention end-to-end: positive score means the metric moves positive - Use the fully linear _NodeGraphToyBridge (Identity ln_final, softmax-free attention projection, activation-free MLP, linear metric) so the first-order identity is exact * feat(attribution_patching): export API + module docs - Export attribution_patch, EdgeAttributionConfig, AttributionResult, and Node from the analysis subpackage __init__ and add them to alphabetized __all__ - Add attribution_patching to the subpackage Tools docstring list - Finalize the module docstring with the algorithm summary, pinned denoising sign/direction convention, names-filter memory note, and node-granularity-only scope - Document that edge scoring (EAP), integrated gradients (EAP-IG), and faithfulness are deferred to follow-on PRs * fix(attribution_patching): capture gradients via bwd hooks, drop metric.backward/zero_grad - Register a backward hook at each cached point instead of retain_grad()/.grad reads, so converted hook points (attn.hook_z, hook_q/k/v, hook_attn_out) yield real canonical-shape gradients rather than None on every real bridge. - Drive the backward with torch.autograd.grad instead of metric.backward() + zero_grad(set_to_none=True), so no caller parameter .grad buffer is clobbered and no model-sized grad buffer is allocated. - Backfill the most-upstream cached point (hook_embed), whose own backward hook cannot fire, from the autograd.grad return; its cached tensor is on-path and unconverted so the returned gradient is exact. * perf(attribution_patching): make the clean forward activation-only - Add compute_gradient flag to cache_activation_and_gradient; when False, register forward hooks only, skip bwd hooks and torch.autograd.grad, and return a cache whose gradients are all None - Run the clean pass with compute_gradient=False so a pair costs two forwards + one backward (as documented) instead of four backward passes; _node_effects reads only clean activations + corrupt gradients * fix(attribution_patching): default names_filter to the node hook set - names_filter=None now falls back to _required_hook_names(n_layers) instead of every hook point; a real Bridge's gated hooks (hook_mlp_in, attn.hook_result, split-QKV inputs) raise in add_hook unless their set_use_* flag is on - update cache_activation_and_gradient docstring so None documents the node-granularity hook set and note an explicit filter is required to reach anything outside it - add unit test asserting names_filter=None on the toy bridge caches exactly {hook_embed, per-layer attn.hook_z, hook_mlp_out} * docs(attribution_patching): drop PR/Risk/step cross-references from user-facing text - Replace PR2/PR3/PR5, Risk 1, and proposal-step labels with plain capability descriptions in module, Node, config, and result docstrings, plus the two NotImplementedError strings and the top_edges raise - Update test match= substrings to edge/integrated gradient, rename test_top_edges_not_implemented_until_pr2 to test_top_edges_not_implemented, and reword the two internal-reference comments * test(attribution_patching): cover the denoising convention on a nonlinear model - Add _NonlinearNodeGraphToyBridge (GELU MLP) so the clean and corrupt runs hold different gradients, exposing which run the score reads from — a property the linear toy bridges structurally cannot test. - Assert attribution_patch scores the corrupt-run gradient: results match the corrupt-gradient convention and diverge from the clean-gradient one, and a single-node corrupt->clean patch confirms the denoising sign end to end. * test(attribution_patching): integration test for finite node scores on a real bridge - Boot a real GPT-2 Bridge and run attribution_patch with the default config and no explicit names_filter, exercising the node-hook-set fallback the toy bridge hides. - Assert the sweep scores exactly the node graph, every score is finite, and all three node families (embed, attn_head_out, mlp_out) are present — the attn_head_out family is the one the hook-conversion bug broke.
…1752) * fix(neox): resolve unembedding as lm_head for transformers >= 5.13 - Map the NeoX bridge unembed component to `lm_head` instead of the removed `embed_out`, so `TransformerBridge.boot_transformers` boots GPT-NeoX/Pythia on the pinned transformers 5.13.0 instead of raising `AttributeError: 'GPTNeoXForCausalLM' object has no attribute 'embed_out'`. - Mirror the same rename on the HookedTransformer weight-conversion path so the bridge and legacy systems stay consistent (AGENTS §2). - Found while generalizing Backward Lens to dense-MLP families, which needs pythia-70m as an `out_in` integration model. * test(neox): cover unembedding resolution on the transformers >= 5.13 layout - Update the top-level HF-path assertion to expect the `lm_head` unembed name. - Add a model-free regression test that resolves the unembed component against a module exposing `lm_head` (and not `embed_out`) via `get_remote_component`, guarding against reintroducing the stale name that broke GPT-NeoX/Pythia boot. * fix(neox): resolve unembedding on both embed_out and lm_head layouts CI's locked transformers==5.13.0 still exposes embed_out; the earlier lm_head-only rename only matched later transformers versions (5.15.1 in the local conda env), breaking every pythia/GPT-NeoX boot in CI. Resolve the unembed name at runtime (bridge: ArchitectureAdapter.prepare_model once the real HF module is available; HookedTransformer: hasattr in convert_neox_weights) so both layouts work. * fix(neox): guard component_mapping before unembed lookup - assert component_mapping is not None in prepare_model before indexing "unembed" - narrows ComponentMapping | None so mypy allows the subscript (fixes index error) * fix(neox): correct transformers embed_out→lm_head rename boundary to 5.14.0 - The GPTNeoXForCausalLM embed_out→lm_head rename landed in transformers 5.14.0 (huggingface/transformers#47198), not 5.13; fix the docstring/comment/test references that said 5.13 or ~5.14 - Rename test_unembed_resolves_against_transformers_5_13_lm_head_layout to _5_14_ so it no longer contradicts the locked-5.13.0-still-has-embed_out fallback test - Leave the four "locked 5.13.0 still exposes embed_out" notes unchanged — that pairing is correct * refactor(neox): set unembed name via components accessor in prepare_model - Replace the redundant `assert component_mapping is not None` + `isinstance(unembed, UnembeddingBridge)` guard with a direct `self.components["unembed"].name = "embed_out"` assignment. - The `components` accessor already asserts the mapping is built and `.name` is declared on GeneralizedComponent, matching how bert.py and other adapters edit their mapping; keep the hasattr `if` guard as-is.
…1754) * test: characterize quantizer-owned scales the itemsize guard misses transformers picks weight_scale_inv's storage dtype from the checkpoint's scale_fmt: one-byte ue8m0 for "ue8m0", float32 for "float", which is the default. Same parameter, same quantizer, same role, two widths. activation_scale is float32 under either format. So cast_floating_params_to_dtype's itemsize < 2 branch protects one spelling of a quantizer-owned scale and rewrites the other. Pins that against the real FP8Linear rather than a hand-rolled fake, since the point is what transformers actually does. * fix: clarify quantized dtype ownership and preserve whole-model skip The skip is right, the reason given for it was wrong, and the wrong reason is what invited #1743. cast_floating_params_to_dtype claimed its itemsize < 2 branch skips "quantizer-owned scale parameters". It does not. It skips one-byte floats, which is a subset. transformers' finegrained-FP8 stores weight_scale_inv as float32 whenever scale_fmt is "float" (the default) and keeps activation_scale float32 under every format, and fbgemm-fp8 stores weight_scale and the expert scales as float32. Reading that docstring, dropping the call-site guard and leaning on the dtype check looks safe. It reintroduces #1713 on those checkpoints instead. So the guard stays, and it stays whole-model, which is also the line transformers draws: to(dtype=) raises for bitsandbytes and GPTQ, half()/float() raise for anything quantized, all keyed on one is_quantized flag. from_pretrained is the component responsible for applying the requested dtype to ordinary floating params, and it has already run by the time this helper is reached, so anything still off that dtype afterward may be quantizer-owned. It also releases at the right moment, because HF deletes quantization_config when a load is dequantized. No runtime behavior change. Restores the test to asserting the skip, now with the evidence for why, and adds a real FP8Linear at both scale widths so the float32 side is covered too. * test: cover quantizer ownership hand-off after dequantization Two things worth pinning about the whole-model guard. It releases: HF deletes quantization_config when a checkpoint is loaded dequantized, so quantization_method drops back to None and normalization runs again. Without that, skipping on the config would strand dequantized checkpoints in their load dtype, which is the one way the guard could be genuinely too broad. And it releases because transformers says so, not because we assume it, so the second test fails if that ever changes upstream instead of letting those checkpoints go quiet. * review: correct the dtype-ownership rationale and tighten the fixtures Review on #1754 caught the docstring asserting from_pretrained applies the requested dtype. It does not, necessarily: nine quantizers override it in HfQuantizer.update_dtype and bitsandbytes, which is what I measured against, is not one of them. AWQ downgrades bf16 to fp16 whenever CUDA or XPU is available regardless of placement, fbgemm-FP8 and FP-Quant force bf16. So the skip defers to the load dtype from_pretrained settled, which on a quantized checkpoint is the quantizer's effective dtype, not the caller's request. Same wording at the call site. Also cut the measurement parenthetical and the to()/half()/float() enumeration, which duplicated the PR description and would rot. 23 lines down to 16. Two more claims of the same kind, found by auditing the rest of the prose rather than waiting for them to be reported: activation_scale is not "float32 either way". It only exists under activation_scheme="static" and the default is "dynamic", where it is None. The unconditional version is weight_scale_inv, float32 unless the checkpoint asks for ue8m0, plus fbgemm-FP8's scales. Biases do not follow the compute dtype. fbgemm-FP8 hardcodes a float32 bias while forcing the model to bf16. The real reason a bias is not quantizer-owned is that param_needs_quantization returns False for it, so that is what the comment says now. Fixture fixes: normalise FP8Linear's float32 bias to the target so the scale-preservation test can only fail on quantizer-owned storage (with the guard removed it now fails on activation_scale and weight_scale_inv, not bias); drive postprocess_model instead of remove_quantization_config, since it is the dequantize branch in the former that the design relies on and calling the latter direct passed even with that branch deleted; skip ue8m0 below torch 2.7, which pyproject allows, because _get_ue8m0_dtype raises instead of falling back. No runtime behavior change.
) * test(jacobian_lens): pin ordinary-estimator fit regression fixture - Capture a deterministic small-model JacobianLens.fit as a golden reference before the PR1 drive-loop refactor. - A tiny locally-built GPT-2 with RNG-independent arithmetic weights makes the fitted transport matrices reproducible across environments, so the upcoming extraction of the estimator-independent driver must preserve these numerics. - Also asserts repeat-fit determinism and joint==merged consistency. * refactor(jacobian_lens): extract estimator-independent fit driver - Add a BackwardProvider seam and _ordinary_vjp wrapper so the fit backward step goes through a swappable callable instead of a direct torch.autograd.grad - Route _jacobian_for_prompt through backward_provider; ordinary path passes _ordinary_vjp so the numerics are identical - Add _fit_transport_matrices driver owning the frozen-parameter lifecycle, prompt accumulation, and source-position averaging for reuse by PR3 - Regression fixture still green, mypy clean; no public API change * refactor(jacobian_lens): route JacobianLens.fit through shared driver - Rewrite fit to delegate to _fit_transport_matrices with the ordinary _ordinary_vjp provider - Remove the duplicated drive-loop body (frozen-parameter lifecycle, per-prompt accumulation, empty-fit guard) * test(jacobian_lens): driver accepts alternate backward provider - Add a seam test that drives _fit_transport_matrices with an alternate backward provider scaling the ordinary VJP by an exact power of two, and assert every transport matrix comes out byte-identical to SCALE*baseline - Confirm the ordinary provider through the driver reproduces the pinned golden regression matrices, anchoring the comparison to the fit fixture - Assert the public JacobianLens.fit path is untouched: it still routes through _ordinary_vjp and reproduces the golden matrices - Wrap two over-length lines in the commit-1 regression fixture * chore(jacobian_lens): format, mypy - annotate shared fit kwargs dicts as dict[str, Any] so ** unpacking type-checks under mypy - import typing.Any in the integration test - format + mypy clean on touched files; regression fixture and seam tests green * fix(jacobian_lens): correct stacklevel drift in short-prompt warning - _fit_transport_matrices's short-prompt warning used stacklevel=2, which was correct only while it was called directly; now that fit delegates through it, that value under-counts by one frame and attributes the warning to fit's own body instead of the caller's fit(...) line. - Bump to 3, and add a regression test that pins the attributed frame via a bare subprocess so this suite's jaxtyping/beartype call-wrapping can't mask the one-frame drift. * test(jacobian_lens): pin one bridge forward per prompt in driver seam test - The alternate-backward-provider test asserted only calls > 0 on the seam. That count tracks dim_batch chunking, not forwards, so it cannot tell one forward per prompt from any number of forwards. - Add a bridge forward pre-hook that counts top-level model(...) calls across the fit and assert it equals len(REGRESSION_PROMPTS): exactly one bridge forward per contributing prompt. - The existing calls > 0 seam check is kept as a weaker guard.
* Bug fix follow ups * MoE fix * CI fixes
* fix: invalidate JacobianLens dictionary cache on weight changes * Optimize JacobianLens dictionary cache validation
* Add clean-coordinate swap clamping * Reuse shared intervention diagnostics * Address Jacobian lens clamp review * Sort Jacobian lens imports * Widen clamp cache type and re-execute Jacobian lens demo - accept either an ActivationCache or the plain dict from run_with_cache(return_cache_object=False) in swap_clamp_hooks, and say so in the docstring - add a unit test that passes the raw cache dict through the clamp - move the demo's rank cell onto swap_clamp_hooks via a shared clamped_logits helper, reword the sweep intro so the live swap is no longer called a clamp, drop the stale swap name from the cleanup cell - re-execute the notebook end to end with the pinned model and lens revisions so the stored outputs match the code, and update the section 3 and 4 prose to this run's numbers --------- Co-authored-by: jlarson4 <jonahalarson@comcast.net>
* test(svd_circuits): add analytic SVD oracle and degeneracy-guard unit tests Add model-free unit tests for a per-head QK/OV singular-vector decomposition. Synthetic weight tensors only, so no model is loaded and no pretrained weights are downloaded. Coverage: - Oracle parity of the factored SVD against torch.linalg.svd for the OV (W_V @ W_O) and QK (W_Q @ W_K.T) maps, plus reconstruction, descending order, the d_head rank bound, and the sigma-ratio rank report. - The .V-not-.Vh convention, asserted by treating the FactoredMatrix.Vh DeprecationWarning as an error. - The degeneracy guard: equal and near-equal singular values group into a block, require_isolated refuses per-direction attribution inside a block and points at the subspace, well-separated spectra flag nothing, eps is honored, near-null runs group, and the blocks partition every direction exactly once. Written test-first: the suite fails at collection until the decomposition module lands. * feat(svd_circuits): add decompose_head with FactoredMatrix QK/OV SVD Decompose a single attention head's query-key (W_Q W_K^T) and output-value (W_V W_O) maps via FactoredMatrix.svd(). Weight-space only: no forward pass, activation cache, or compatibility mode. Each map stays factored, so the d_model x d_model product is never materialized and the rank is bounded by d_head. Right singular vectors are read from FactoredMatrix.V; the deprecated .Vh alias is never touched. Each direction is summarized in a rank report with its singular value, its ratio to the top value, and the contiguous block it groups into, so near-equal singular values surface as rotation-ambiguous subspaces. The module is importable by full path only and is intentionally absent from the analysis package __all__. * feat(svd_circuits): add degeneracy guard refusing per-direction attribution Add HeadSVD.is_degenerate, block_of, degenerate_blocks, and require_isolated so callers can detect rotation-ambiguous singular blocks and refuse per-direction attribution inside them. Near-equal singular values leave the singular directions defined only up to a rotation within their block, so require_isolated raises DegenerateDirectionError and points at the block as a subspace instead of returning a rotation-dependent direction. The methods read the existing per-direction report (block_id and is_degenerate are already populated by _degeneracy_blocks and _build_rank_report), so no grouping logic is recomputed. * fix(svd_circuits): read one layer's weights and map query heads to key-value heads - Read the attention block's own per-head W_Q/W_K/W_V/W_O instead of the full-model W_Q/W_K/W_V/W_O stacks, so decompose_head materializes one layer instead of the whole model's weights. - On grouped-query attention W_K/W_V carry one row per key-value head, so a query head index was out of range or wrong once key-value heads outnumbered query heads; map the query head to its key-value head (h // (n_heads // n_kv_heads)) before indexing, mirroring the bridge's own kv-head expansion rule. - A no-op for multi-head attention, where the head counts already match. - Add model-free unit tests exercising decompose_head directly through a lightweight stub model (no from_pretrained, no download): head selection under multi-head attention, the query-to-kv-head mapping under grouped-query attention, the QK transpose wiring, the which filter, and the layer/head/which input guards. * fix(svd_circuits): bound degenerate blocks by distance to the block anchor - Require a new direction to fall within eps of both its immediate predecessor and the block's anchor (its first, largest member) before it joins a near-equal block, capping how far a block can spread. - Without the anchor bound, a spectrum decaying by just under eps at every step chained every direction into one block whose extremes differed by far more than eps, so require_isolated wrongly refused the top direction over the whole spectrum. - Add a regression test for a slowly decaying spectrum: it no longer collapses into one block, and the well-separated top direction stays isolated. - Add a test pinning the exact block degenerate_blocks() returns for a repeated singular value, rather than only checking non-degeneracy. * fix(svd_circuits): base the null-space cutoff on effective rank - Add a dedicated null_rtol parameter to decompose_head, threaded through _factored_head_svd into _degeneracy_blocks, so the null-run grouping no longer reuses the near-equal eps threshold - Default null_rtol to d_model * torch.finfo(S.dtype).eps, matching torch.linalg.matrix_rank's default relative tolerance, so directions group as null only when they reflect the map's numerical rank - Store the resolved null_rtol on HeadSVD for reproducibility - Update _degeneracy_blocks and the module docstring to describe the null cutoff as distinct from the near-equal gap threshold eps - Rewrite test_null_run_grouped to use a spectrum the relative-gap rule cannot group, grouped only via an explicit null_rtol - Add test_null_default_does_not_overgroup and test_null_run_groups_true_null_tail to pin the default cutoff's behavior on both a spectrum it should not over-group and a genuinely null tail it should still catch * docs(svd_circuits): rewrite the example on TransformerBridge and drop HookedTransformer - Rewrite the module docstring Example:: to boot a TransformerBridge instead of HookedTransformer, matching the per-block accessor path decompose_head already reads. - Drop the "Works with both HookedTransformer and TransformerBridge" sentence and the HookedTransformer mention in the Args: block; state the per-block bridge accessors decompose_head reads instead. - Docstring-only: no HookedTransformer-only code path existed to remove. * Dropped the anchor bound on degeneracy blocks, flag a lone null direction as degenerate, read per-head weights detached, widen which to Sequence[str], adjust tests to match above --------- Co-authored-by: jlarson4 <jonahalarson@comcast.net>
…h, and stop mutating the ActivationCache (#1745) * fix(head_detector): forward kwargs on the multi-prompt path and stop mutating the cache detect_head recursed with positional arguments only when given a list of prompts: batch_scores = [detect_head(model, seq, detection_pattern) for seq in seq] so error_measure, exclude_bos, exclude_current_token and heads were all silently reset to their defaults. Callers asking for error_measure=abs received mul scores instead, on a different scale and sometimes with the opposite sign; a heads filter was ignored entirely. The repo's own Head_Detector_Demo notebook calls detect_head on a list with error_measure=abs, so its published plots show mul scores. Separately, compute_head_attention_similarity_score masked exclude_bos and exclude_current_token in place. attention_pattern is a view into the caller's ActivationCache, so the cached pattern was permanently modified: rows stopped summing to 1 and every later read of that cache saw altered activations. Clone before masking. Adds seven regression tests: one per forwarded keyword argument, and three asserting the cache is untouched and its rows still sum to 1. * fix(head_detector): reject cache on the multi-prompt path instead of dropping it Review follow-up. The multi-prompt path now forwards heads, exclude_bos, exclude_current_token and error_measure, but cache was still being dropped without a word. Forwarding it is not an option: an ActivationCache holds the activations of one prompt, so reusing it across a list would score every prompt against the first one's attention patterns. Silently ignoring it is the same failure this PR exists to fix, one argument along. So the list path rejects it explicitly. Adds three tests: passing a cache with a list raises, passing a cache with a single prompt still matches the uncached result, and the list path without a cache is unchanged. * fixing formatting --------- Co-authored-by: jlarson4 <jonahalarson@comcast.net>
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
6 participants
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.
Description
Pre-release master commit for version 3.9.0
New tools
BackwardLensprojects the low-rank factors of GPT-2 MLP weight gradients into vocabulary space (Katz et al., 2024), turning one backward pass into per-layer token readouts of what each MLP matrix is being pushed to learn. GPT-2 only for now; ships with a demo notebook and docs page. (feat(backward_lens): add vocabulary readout of GPT-2 MLP gradient factors #1723)attribution_patchscores every residual-stream node on aTransformerBridgefrom one clean forward and one corrupt forward/backward, instead of one forward per intervention. Node-level scores in this release; edge attribution (EAP) and EAP-IG are planned follow-ups. (feat(attribution_patching): node attribution-patching substrate on TransformerBridge #1750)temporarily_swap_parameter– context manager intransformer_lens.utilitiesthat swaps a parameter and restores it in both the bridge and the wrapped HF model, so temporary edits can't become permanent. (Creating a parameter swap utility #1744)Jacobian Lens
alpha) while holding the residual and every other coordinate fixed. Offline viaJacobianLens.coordinate_patch(feat(jacobian_lens): anchored J-space coordinate patching (offline primitive) #1741), or live insiderun_with_hooks/generateviaJacobianLens.coordinate_patch_hooks(feat(jacobian_lens): add dynamic J-space coordinate-patch hooks #1749).JacobianLens.swap_clamp_hookspins each hooked layer's swapped coordinates to their clean-run values; the docs now note that repeated liveswap_hooksswaps can cancel across a layer band. (Add clean-coordinate clamping for JacobianLens swaps #1747)lens_vector_dictionary()now invalidates its cache whenW_Uchanges (optimizer step, in-place edit, parameter swap, a different model) instead of returning directions under the old token labels. (Fix JacobianLens dictionary cache invalidation #1766)fit()requires eval mode – fitting now raises if any module is in training mode, since dropout made the fitted Jacobian depend ondim_batch. (Require eval mode for JacobianLens fitting #1763)Generation
generate_stream(return_type="str")decodes every batch row instead of only the first (fix(bridge): decode every batch row in generate_stream #1757), and no longer corrupts multi-byte characters (CJK, emoji) that land on a chunk boundary (Followup review #1770).hf_generate– batched string prompts are left-padded and the attention mask is forwarded, so shorter prompts no longer generate from pad tokens. (Fix batched hf_generate padding and attention masks #1761)generate()masks and padding – the attention mask now reaches stateful/SSM models (Mamba, Bamba, Nemotron-H), and a failed tokenization no longer leaves the shared tokenizer stuck on left-padding. (Followup review #1770)Models and weight processing
embed_out→lm_head, which broke every NeoX checkpoint inboot_transformers; the bridge adapter and HookedTransformer conversion are both updated. (fix(neox): resolve unembedding as lm_head for transformers >= 5.13 #1752)process_weights()call now warns and skips instead of re-applying the fold (Followup review #1770).LNPre/RMSPreforboot_nativemodels that must not carry norm weights (SoLU, attn-only, Othello-GPT). (Followup review #1770)-infattention masks – BLOOM and GLM-DSA attention scores report masked positions as-infrather thanfinfo.minin compatibility mode, sotorch.isinfmasking works. (Followup review #1770)HookedAudioEncoderzeroes padded frames before the positional convolution, matching HF, so padding no longer bleeds into real frames (Fix Legacy Audio System #1736).Bridge core
original_modelno longer duplicated –.to()/.cuda()had registered the HF model as a submodule, doubling it instate_dict()and moving its weights twice;original_modelis now read-only. (Followup review #1770)use_attn_resulton a deepcopied bridge no longer rewires the original's attention. (Followup review #1770)str(bridge)works again – it raisedAttributeError. (Followup review #1770)get_bridge_paramsno longer returns an all-zeropos_embed.W_pos(up to ~2 GB) for rotary models. (Followup review #1770)Head detector
detect_headon a list of prompts now honorserror_measure,exclude_bos,exclude_current_token, andheads(all were silently reset to defaults), and no longer zeroes attention patterns inside the caller'sActivationCache. (fix(head_detector): forward keyword arguments on the multi-prompt path, and stop mutating the ActivationCache #1745)Upgrade notes
jaxtyping>=0.3is now required. (CI Warnings Cleanup #1732)JacobianLens.fit()raises on a model in training mode; callmodel.eval()first. (Require eval mode for JacobianLens fitting #1763)generate_stream(..., return_type="str")yieldsList[str]rather than the first row'sstr. (fix(bridge): decode every batch row in generate_stream #1757)bridge.original_modelcan no longer be assigned. (Followup review #1770)detect_head(model, [prompts...], cache=...)now raisesValueErrorinstead of silently ignoring the cache. (fix(head_detector): forward keyword arguments on the multi-prompt path, and stop mutating the ActivationCache #1745)Also in this release
svd_circuits.decompose_head, not yet exported) (feat(svd_circuits): per-head QK/OV SVD with degeneracy guard #1768) and for a relevance lens (refactor(jacobian_lens): extract estimator-independent fit driver #1764); BERT bridge-vs-HF parity tests (test(bert): add Bridge-to-HF parity coverage #1735); quantized-dtype ownership docs (Clarify quantized dtype ownership; keep the whole-model skip (#1743) #1754); CI warning cleanup (CI Warnings Cleanup #1732).Type of change
Please delete options that are not relevant.
Checklist: