Models and caching

Operations accept any of three model types:

  • a TransformerBridge

  • a HookedTransformer

  • either of them wrapped in a CachedModel

TransformerBridge

TransformerBridge is how transformer-lens 3 and later load a HuggingFace model. entoscope reads it as it is, so its activations and weights are HuggingFace’s own, and a run means the same as the same computation done on the HuggingFace model directly:

from transformer_lens.model_bridge import TransformerBridge

model = TransformerBridge.boot_transformers("gpt2")
model.eval()

boot_transformers hands back a model whose HuggingFace weights are already in eval mode, so the eval() is a safeguard. It matters for a HuggingFace model you build yourself and pass in as hf_model=: the bridge’s eval() does not reach the model it wraps, so call eval() on that model too, or its dropout stays on.

Compatibility mode

model.enable_compatibility_mode() switches a bridge to transformer_lens’s own conventions. It folds each LayerNorm into the weights that follow it, and centers the weights that write to the residual stream. Both modes work with every operation. What changes is what the intermediate numbers mean, not whether a run works.

One consequence shows in the parameter lens. GPT-2 ties its embedding and unembedding, so with HuggingFace’s weights W_E.T equals W_U and the lens’s two projection spaces agree exactly; folding and centering separate them.

HookedTransformer

HookedTransformer is accepted wherever a bridge is, with the transformer-lens releases that have it: 2.x and 3.x, since transformer-lens 4 removed it in favour of the bridge. It is the only option on Intel macOS, whose transformer-lens 2.15 has no bridge:

from transformer_lens import HookedTransformer

model = HookedTransformer.from_pretrained("gpt2")
model.eval()

from_pretrained folds LayerNorm and centers weights by default, so its numbers match a bridge in compatibility mode rather than a plain bridge.

Caching activations on disk

CachedModel wraps a model and stores what run_with_cache returns on disk, so a forward pass a run has already made is read back instead of recomputed. Everything else is forwarded to the wrapped model, so an operation cannot tell the difference:

from entoscope import CachedModel

cached = CachedModel(TransformerBridge.boot_transformers("gpt2"), cache_dir=".activation_cache")
cached.eval()

With cache_dir=None the wrapper passes every call straight through, which lets a script make caching optional with a single argument. A method that returns the model itself, such as to or eval, returns the wrapper, so cached = cached.to("cuda") keeps the cache, and the wrapper is called as the model is.

What gets cached

Only run_with_cache calls are cached, which are the runs an input makes before the method starts:

  • the lens’s prompt

  • every probing example

  • patching’s source and target runs

  • every circuit-discovery prompt, with the residual stream, MLP output, values and attention pattern of every layer

Patched runs and gradient passes go through run_with_hooks and are never cached. Circuit discovery caches those activations of every layer per prompt, which takes far more disk than an index-filtered run.

What an entry is keyed by

An entry is stored at <cache_dir>/<model name>/<digest>.pt, written atomically. Its digest covers:

  • the tokens, the requested activation names in any order, and every other run_with_cache argument that changes the result;

  • a fingerprint of the model: its class, its config’s dimensions and dtype, and the name, shape, dtype and first 64 values of every parameter;

  • the version of the entry format, so an entry written by a release that stored it differently is recomputed rather than misread.

The weights in the fingerprint keep a fine-tune, a supplied hf_model or another revision from sharing entries with the base model. A bridge and a HookedTransformer of the same model never share entries either.

The fingerprint is taken on every call, so a change made after wrapping, such as enable_compatibility_mode(), a change of dtype or a training step, is a new entry rather than a stale one. Patching runs each prompt as wide as its patched runs, and a tensor whose rows come out alike bit for bit is stored as one row. A backend that rounds the copies apart, as some x86 CPUs do, has every row kept, so an entry always reads back exactly what the model computed. Entries are never evicted; delete the directory to clear the cache.