Conversation
* Fixed issue 1647 * OLMo fixes * formatting
* Fixed issue 1647 * OLMo fixes * Granite resolution * comment cleanup
* initial numerics fix * Additional clarification and bug cleanup
* fix DeepSeek V2 dense MLP hooks * Isolate DeepSeek compatibility fixtures
) * fix(bridge): respect use_parallel_residual in the GPTNeoX adapter The NeoX adapter hardcoded the parallel-residual wiring, ignoring HF's use_parallel_residual. On the sequential branch HF computes a genuine post-attention residual that post_attention_layernorm reads, but ParallelBlockBridge pops the hook_resid_mid alias by design -- so the hook was silently missing on exactly the checkpoints that have one, and cfg.parallel_attn_mlp reported True for a model that is not parallel. Logits stayed correct because the bridge delegates to HF's own block, so nothing raised. All six RedPajama-INCITE checkpoints in the registry are GPTNeoXForCausalLM with use_parallel_residual=false, so resid-mid patching, attribution and SAE work were unavailable on registered models. Pythia is genuinely parallel and is unaffected. Select the block class from the flag, mirroring the guards already in stablelm.py:121 and falcon.py:148, and add use_parallel_residual to _HF_PASSTHROUGH_ATTRS so the adapter can see it -- reading the resolved cfg.parallel_attn_mlp instead would flip hand-built configs to sequential, since it defaults to False while HF's NeoX default is True. Mirror the same hardcode on the HookedTransformer side. Add test_parallel_residual_identities.py asserting the identity ParallelBlockBridge documents in its own docstring (resid_post == resid_pre + attn_out + mlp_out) across all seven adapters that use it, exercising the three config-switchable families in both wirings. Fixtures are seeded tinies from local HF configs, so no hub access is needed. Fixes #1644 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(bridge): honour parallel_attn_mlp on caller-supplied NeoX configs Review feedback from @jlarson4 on #1649. TransformerBridgeConfig has no use_parallel_residual field, so on the build_bridge_from_module(..., tl_config=...) path the getattr default fired and a config with parallel_attn_mlp=False still got ParallelBlockBridge with no hook_resid_mid. Fall back to parallel_attn_mlp before defaulting to HF's True. The existing _make_cfg fixture relied on the old hardcode, so it now states parallel_attn_mlp=True explicitly -- a caller-supplied NeoX config is otherwise indistinguishable from one that asked for sequential, since the dataclass default is False. Wrap the tiny-model seeding in test_parallel_residual_identities.py in torch.random.fork_rng(devices=[]), matching bridge.py:360. Seeding ran only on a cache miss, so it leaked RNG state into later unseeded tests depending on ordering. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * test(bridge): pin use_parallel_residual passthrough The fallback added for caller-supplied configs made the _HF_PASSTHROUGH_ATTRS entry redundant -- reverting it left every test green, since sources/transformers.py already derives parallel_attn_mlp and the fallback picks it up. That only holds while TransformerBridgeConfig keeps defaulting parallel_attn_mlp to False, so assert the adapter sees HF's own flag rather than resting on two defaults agreeing. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
…#1666) * initial numerics fix * Additional clarification and bug cleanup * Fixing issues with MoE and Dense hooks * test cleanup and improvements
* initial numerics fix * Additional clarification and bug cleanup * Fixing issues with MoE and Dense hooks * test cleanup and improvements * Mixtral 5.x bug fix
…g them as matrices (#1669) * initial numerics fix * Additional clarification and bug cleanup * Fixing issues with MoE and Dense hooks * test cleanup and improvements * Mixtral 5.x bug fix * Quantization bugs * Missing tests restored * Cleanup of potential errors * fix formatting
* feat(jacobian_lens): add J-space sparse-decomposition solver - Add `get_sparse_decomposition` to decompose an activation into a k-sparse nonnegative combination of J-lens vectors (Gurnee et al., 2026). - Support `nonnegative_orthogonal_matching_pursuit` (default, exact NNLS re-solve) and `gradient_pursuit` algorithms. - Return both the nonnegative coordinates and the orthogonal-projection J-space component. - Distinguish the projection from the coefficient reconstruction; the projection residual matches `swap_hooks`. - Keep the implementation model-free by operating directly on the raw dictionary tensor. - Add tests covering both algorithms, exact-resolve NNLS correctness, a brute-force optimum oracle, and input validation. Part of #1539 (Tier 2). * feat(jacobian_lens): add full-vocabulary lens-vector dictionary - Add `JacobianLens.lens_vector_dictionary(model, layer)` returning the `[d_vocab, d_model]` dictionary whose rows are the J-lens vectors `v_t = J[layer]^T W_U[:, t]`. - Cache the dictionary per (layer, device) and release it in `clear_device_cache`, so a sparse decomposition can reuse it; document its vocabulary-sized memory cost. - Add tests asserting the dictionary matches `lens_vectors` over every token, is cached and invalidated by `clear_device_cache`, and rejects an unfitted layer. * feat(jacobian_lens): add JacobianLens.decompose wrapper and exports - Add `JacobianLens.decompose(model, activation_or_prompt, layer, *, position, k, algorithm)` decomposing either a raw activation vector or the `blocks.{layer}.hook_out` activation at a prompt position, validating inputs before building the dictionary. - Build and cache the layer dictionary via `lens_vector_dictionary` and solve with `get_sparse_decomposition`. - Export `JSpaceDecomposition` and `get_sparse_decomposition` from `transformer_lens.tools.analysis`. - Add end-to-end tests for the raw-activation and prompt paths, the algorithm passthrough, and the input-validation error paths. * test(jacobian_lens): add real-model decompose tests and docs - Add a GPT-2 integration test (regular CI): `decompose` on a real `blocks.6.hook_out` activation returns k nonnegative atoms, the non-J-space residual is orthogonal to every selected J-lens vector, and the J-space component plus residual recover the activation. - Add a slow gemma-2-2b-it integration test validating `decompose` on the published lens artifact: support size, nonnegative coordinates, in-vocabulary token ids, and component-plus-residual reconstruction. - Document J-space sparse decomposition in `jacobian_lens_fitting.md`: the `decompose` API, local coordinates versus the orthogonal-projection J-space component, and the paper's variance facts with closed-model caveats. * docs(jacobian_lens): cite the decomposition algorithm sources Add a References section to the decomposition module docstring: Gurnee et al. (2026) for the J-space method, Pati et al. (1993) for the greedy orthogonal-matching-pursuit selection, Blumensath & Davies (2008) for the gradient-pursuit update, and Lawson & Hanson (1974) for the active-set nonnegative least-squares re-solve. * fix(jacobian_lens): complete and independently validate NNLS Resolve PR #1596 review comments 1 and 3 as one numerical-correctness unit: the drop-only active-set approximation could strand an atom that is optimal later, returning a non-KKT point (the reviewer measured 4/288 GPT-2 decompositions off, relative dual violation up to 0.55). Solver (`_nonnegative_least_squares`): - Replace the drop-only loop with the full Lawson-Hanson active-set method, so a released atom can re-enter. Solve the passive set in float64 with an explicit pseudoinverse rank threshold. - Fail closed: the feasibility corrections use the classical `3 * num_active` budget with a separate admission safeguard; an exhausted budget, an invalid line-search step, a stalled correction, or a failed KKT check raises `RuntimeError` rather than clamping and returning an unverified vector. - Enforce the passive-set invariant and guard the blocking-ratio step against zero denominators and zero-current/zero-candidate degeneracy. - Use one scale- and dtype-aware tolerance policy (`_nnls_tolerances`) for dual feasibility, coefficient cleanup, and the KKT test. - Validate the KKT conditions (`_validate_nnls_kkt`) before returning, in both the float64 work dtype and the caller's result dtype. Independent validation (tests): - Add `_reference_nnls`, a brute-force support-enumeration NNLS that shares no code with the solver, and compare objectives *two-sidedly* over many shapes. - Add an independent `_assert_nnls_kkt` certificate (primal/dual feasibility, stationarity, complementarity) with a deliberately looser, scale-aware tolerance, used on rank-deficient, duplicate-column, near-collinear, boundary, zero-target, and jointly-rescaled systems. - Add a fail-closed test proving the safeguard raises instead of returning an unverified vector, plus dtype/device, zero-row-invariance, and realistic width (768x25) checks. Document the new public contract: `decompose` now raises `RuntimeError` on a KKT-uncertifiable solve, and the fitting docs describe the float64 re-solve and KKT check. Verification: - tests/unit/tools/test_jacobian_lens_decomposition.py: 127 passed. - KKT reproduction over 935 real GPT-2 decompositions on the full branch (commits 6-7; the NNLS solver added here is unchanged by commit 7): 0 KKT failures, max relative dual violation 2.77e-08 (old drop-only solver: up to 0.55 on 4/288). - black/isort/pycln clean; mypy clean on the touched source. * fix(jacobian_lens): distinguish active and selected support Resolve PR #1596 review comment 2: the greedy loop ran exactly `k` times and kept every selected atom in `support`, so a coordinate the NNLS solve drove to zero still consumed a support slot and was never reconsidered. On GPT-2 this returned `support` of size 25 with only ~9 nonzero coordinates ("dead slots"), and the docs presented all 25 as active concepts. Make `k` an upper bound and separate the two supports the paper conflates: - Early stopping: selection stops once no unselected atom is materially positively correlated with the residual (under nonnegativity a negatively-correlated atom cannot reduce it). The stop threshold sits at the float32 residual noise floor, so a full-rank target stops instead of selecting noise atoms -- this makes the selected support scale-invariant. - `support` is now the numerically *active* set: selected atoms whose contribution `c_i * ||v_i||` is a materially nonzero fraction of `||x||` (a scale-invariant activity threshold matching the NNLS coefficient-zeroing scale, so `support` equals the strictly-positive NNLS coordinates). `coordinates` is aligned with it and every entry is strictly positive. - New `selected_support` holds every greedily selected atom and defines the span for `j_space_component`. Hence `support <= selected_support <= k`. - `reconstruction` is the nonnegative combination over the active support; for the exact NNLS re-solve it equals the projection onto that support (KKT stationarity), so it differs from `j_space_component` exactly when a selected atom has a zero coordinate. Empty sets (zero target, or a target orthogonal to every atom) are handled explicitly. - `gradient_pursuit`: its projected line-search step uses a backtracking line search -- the exact unconstrained step is projected onto the nonnegative orthant, then halved until the projected update no longer increases the residual (falling back to the feasible incoming point if the bounded search finds none), so the objective is monotonically non-increasing. Its `support` is filtered to the final active coordinates like the default algorithm. - Autograd: the scalars used only for control flow (the target norm, the early-stop correlation gate, and the gradient-pursuit residual comparison) are detached, so the model-free primitive never pulls caller-owned tensors out of the graph; the autograd-contract test backpropagates through the reconstruction and asserts finite gradients on both inputs. - Keep `support`/token tensors on CPU and vector outputs on the compute device. Propagate the contract through the wrapper, tests, and docs: `decompose` and `jacobian_lens_fitting.md` document `k` as an upper bound and the active vs selected/projection distinction, naming which operationalization the paper's variance figures measure. Integration and unit assertions check `support <= selected_support <= k`, subset, and all-active instead of `support == k`. Verification (conda env, python 3.12, H100): - tests/unit/tools/test_jacobian_lens_decomposition.py + test_jacobian_lens.py: 219 passed. - tests/integration/test_jacobian_lens.py (GPT-2): 13 passed. - mypy clean on the touched source. The only `mypy .` error is pre-existing and unrelated (olmo_hybrid.py, from transformers 5.15.0 vs the locked 5.13.0). - black/isort/pycln clean on the touched files. - KKT reproduction over 935 real GPT-2 decompositions (3 prompts x 11 source layers x all positions, k=25): 0 dead slots (support == strictly-positive coordinates), 322/935 with selected > active, 0 KKT failures, max relative dual violation 2.77e-08 (old drop-only solver: up to 0.55). - Scale-invariance of support verified across many seeds; CUDA device-safety checked. --------- Co-authored-by: Jonah Larson <jonahalarson@comcast.net>
* Hook extensions for specialized bridges * Additional bug fixes
* Hook extensions for specialized bridges * Additional bug fixes * Comment cleanup
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
4 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
A series of Bug fixes aimed at improving compatibility with
HookedTransformer, and expanding overall hook coverageType of change
Checklist: