Fix JacobianLens dictionary cache invalidation

[emerardd]

Description

JacobianLens.lens_vector_dictionary() previously cached a full-vocabulary dictionary only by layer and device. Reusing the lens after an optimizer step, an in-place or temporary unembedding edit, a parameter replacement, or with another compatible model could therefore return a stale dictionary and silently attach residual directions to the wrong token labels.

This change retains one detached W_U snapshot per device and compares it with the current unembedding using on-device torch.equal. Unchanged weights retain the existing object-level dictionary cache hit; a content change invalidates every cached layer on that device. This covers .data writes that bypass PyTorch's version counter and parameter replacements after allocator address reuse, without copying weights to the host or hashing them on every cache lookup.

The regression coverage verifies cache hits for unchanged weights and invalidation after copy_(), an optimizer step, .data mutation without a version increment, whole-parameter replacement after the old parameter has been released, temporary parameter swapping in both directions, and reuse with a second compatible model. It also verifies that one snapshot is shared across layers, a change invalidates every layer on the affected device, and an unembedding row swap cannot make decomposition return the old token label with zero residual.

Fixes #1765

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Screenshots

Not applicable.

Validation

• uv run --no-cache --no-sync pytest tests/unit/tools/test_jacobian_lens.py tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py -q — 122 passed
• uv run --no-cache --no-sync mypy . — success across 397 source files
• pycln, isort, Black, and git diff --check — passed for the affected surface

The complete make test-pr surface was not run locally.

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 — only the affected unit-test surface was run
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

[jlarson4]

@emerardd Great work on this, thanks for bringing this issue to our attention. Just a couple small comments on the _unembedding_fingerprint function

[jlarson4]

.data deliberately does not share a tensor's version counter, and a .data write also leaves data_ptr, shape, stride, dtype, device and the model object identical. Every fingerprint field matches and the stale dictionary is served. #1765's headline symptom survives here: swap unembedding rows with weight.data[[0, 1]] = weight.data[[1, 0]] and decompose labels the current token-0 direction as token 1 with a zero residual. weight.data.copy_() and weight.data.fill_() are an established pattern in this repo — test_linear_bridge.py uses both.

Please base the fingerprint on the unembedding's contents rather than on its identity metadata; a cheap digest costs far less than the d_model² × d_vocab matmul it guards.

[emerardd]

Fixed in c84e4a7. The fingerprint now hashes the logical W_U contents with BLAKE2b (plus shape, dtype, and device), so .data writes no longer depend on _version. I added the row-swap regression and assert the version remains unchanged; decomposition selects the current token label with zero residual.

[jlarson4]

A small follow up on this:

Would it be possible to detect the change on-device instead of hashing host bytes? Something like torch.equal against a retained copy of W_U.

As it's built here, .contiguous() always materializes a full copy, .cpu() adds a device->host transfer and a sync on a GPU model, and then blake2b hashes the whole buffer single-threaded, against a matmul that is multi-threaded BLAS. At GPT-2 scale the guard costs roughly eight times the rebuild it protects, so the cache is now slower than having no cache at all, and it is still a net loss at Llama-3-8B's d_model.

[emerardd]

Addressed in 5d76b4c. I replaced the host-side BLAKE2b fingerprint with one detached W_U snapshot per device and compare it with the current unembedding using on-device torch.equal. When the contents change, all dictionary entries on that device are invalidated and the single shared snapshot is refreshed, avoiding both the device-to-host transfer and per-layer snapshot duplication. clear_device_cache() now releases the snapshots as well. I added coverage that verifies the snapshot is shared across layers and that a mutation invalidates every cached layer on the affected device.

[jlarson4]

Awesome, thank you @emerardd, looks great, merging now

[koriyoshi2041]

Validated the latest on-device snapshot implementation at 5d76b4c6: the two JacobianLens test modules pass 122/122 locally, including .data mutation, parameter replacement, cross-model reuse, and cross-layer invalidation; diff hygiene is clean. The shared per-device snapshot also removes the earlier host-transfer/hash cost while preserving content-based invalidation.