Info

Last Execution: 2026-08-28

Package	Version
nnsight	0.7.1.dev41+gd901da3ed
Python	3.12.13
torch	2.13.0+cu126
transformers	5.15.0.dev0

Migrating from TransformerLens¶

You have a TransformerLens script and you want to run it in nnsight. This page is the map — and, more importantly, the warning about the one thing that will make a correct port produce wrong numbers.

The short version. nnsight traces a HuggingFace model as it is. TransformerLens does not: HookedTransformer.from_pretrained rewrites the weights into an equivalent-but-different parameterisation. So some quantities are identical between the two libraries and some are legitimately different, and knowing which is which is the whole difficulty of a port.

	agree to ~1e-5	differ by convention
Behavioural — logit differences, patching effects	✅ (identical to 4 d.p. below)
Decompositional — per-head DLA, anything reading `z`, `v`, or `resid` directly		⚠️

If your ported patching map matches and your direct logit attribution doesn't, you have found a convention, not a bug.

Everything below was run on GPT-2 small against a live TransformerLens cache; the diffs printed are real output from the cells you can see.

Setup¶

We load the same checkpoint twice — once through each library.

In [1]:

Copied!





import time, torch, nnsight
from transformer_lens.model_bridge import TransformerBridge
from nnsight.modeling.transformers import TransformersModel

torch.set_grad_enabled(False)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# --- TransformerLens: weight processing ON (the default people use) ---
tl = TransformerBridge.boot_transformers("gpt2")
tl.enable_compatibility_mode(
    center_unembed=True, center_writing_weights=True, fold_ln=True,
)   # refactor_factored_attn_matrices=True would additionally rewrite Q/K and the
    # OV circuit, after which hook_q / hook_k / hook_z are no longer comparable
tl.set_use_attn_result(True)

# --- nnsight: the raw HuggingFace model, no processing ---
# attn_implementation="eager" is required to reach attention probabilities.
nn = TransformersModel("openai-community/gpt2", task="text-generation",
                       dispatch=True, device_map=DEVICE,
                       attn_implementation="eager")

PROMPT = "When John and Mary went to the shops, John gave the bag to"
import time, torch, nnsight
from transformer_lens.model_bridge import TransformerBridge
from nnsight.modeling.transformers import TransformersModel

torch.set_grad_enabled(False)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# --- TransformerLens: weight processing ON (the default people use) ---
tl = TransformerBridge.boot_transformers("gpt2")
tl.enable_compatibility_mode(
    center_unembed=True, center_writing_weights=True, fold_ln=True,
)   # refactor_factored_attn_matrices=True would additionally rewrite Q/K and the
    # OV circuit, after which hook_q / hook_k / hook_z are no longer comparable
tl.set_use_attn_result(True)

# --- nnsight: the raw HuggingFace model, no processing ---
# attn_implementation="eager" is required to reach attention probabilities.
nn = TransformersModel("openai-community/gpt2", task="text-generation",
                       dispatch=True, device_map=DEVICE,
                       attn_implementation="eager")

PROMPT = "When John and Mary went to the shops, John gave the bag to"

Gotcha 0: tokenization¶

model.to_tokens() prepends a BOS token. tokenizer(...) does not. Feed the two libraries different tokens and every number after this point disagrees for a reason that has nothing to do with your port — so we tokenize once, with TransformerLens, and hand the same ids to nnsight.

In [2]:

Copied!





toks = tl.to_tokens([PROMPT])                        # prepends BOS
hf_toks = nn.tokenizer(PROMPT, return_tensors="pt")["input_ids"]
print("TL to_tokens :", toks.shape[-1], "tokens")
print("HF tokenizer :", hf_toks.shape[-1], "tokens  <-- no BOS")

toks = toks.to(DEVICE)
B, S = toks.shape
H, D = tl.cfg.n_heads, tl.cfg.d_model
d = tl.cfg.d_head
toks = tl.to_tokens([PROMPT])                        # prepends BOS
hf_toks = nn.tokenizer(PROMPT, return_tensors="pt")["input_ids"]
print("TL to_tokens :", toks.shape[-1], "tokens")
print("HF tokenizer :", hf_toks.shape[-1], "tokens  <-- no BOS")

toks = toks.to(DEVICE)
B, S = toks.shape
H, D = tl.cfg.n_heads, tl.cfg.d_model
d = tl.cfg.d_head

TL to_tokens : 15 tokens
HF tokenizer : 14 tokens  <-- no BOS

1. The four conventions that silently break a port¶

enable_compatibility_mode / from_pretrained apply four transformations. Each one changes what a hook returns, without changing what the model computes.

Flag	What it does	What it breaks
`center_writing_weights`	every matrix writing to the residual stream is made mean-zero over `d_model`	`hook_resid_*` differs from HF `hidden_states` by that position's mean
`center_unembed`	`W_U` made mean-zero over vocab	raw logits differ by a per-position constant — cancels in any logit difference
`fold_ln`	LayerNorm gain folded into the next matrix, bias into its bias	`ln1.hook_normalized` is not `ln_1.output` — it has no gain or bias
`fold_value_biases`	`b_V` pushed forward into `b_O`	`hook_v` and `hook_z` are `v_HF − (ln_1.bias @ W_V + b_V)`

Let's measure all of it rather than take it on faith.

In [3]:

Copied!





_, tl_cache = tl.run_with_cache(toks)

L = 5
with nn.trace(toks):
    resid_pre = nn.transformer.h[L].input.save()            # block input
    z_in      = nn.transformer.h[L].attn.c_proj.input.save()  # what TL calls z
    attn_out  = nn.transformer.h[L].attn.output[0].save()      # tuple! see gotchas

center = lambda x: x - x.mean(-1, keepdim=True)

def report(name, tl_val, nn_val, do_center=True):
    raw = (tl_val - nn_val).abs().max().item()
    line = f"{name:24} raw max|TL-nn| = {raw:8.4f}"
    if do_center:
        cen = (center(tl_val) - center(nn_val)).abs().max().item()
        line += f"   after centering both = {cen:.2e}"
    print(line)

report("resid_pre",  tl_cache["blocks.5.hook_resid_pre"], resid_pre)
report("attn_out",   tl_cache["blocks.5.hook_attn_out"],  attn_out)
report("z",          tl_cache["blocks.5.attn.hook_z"],
       z_in.view(B, S, H, d), do_center=False)
print("   ^ z differs by the value-bias fold; this is expected, not a bug.")
_, tl_cache = tl.run_with_cache(toks)

L = 5
with nn.trace(toks):
    resid_pre = nn.transformer.h[L].input.save()            # block input
    z_in      = nn.transformer.h[L].attn.c_proj.input.save()  # what TL calls z
    attn_out  = nn.transformer.h[L].attn.output[0].save()      # tuple! see gotchas

center = lambda x: x - x.mean(-1, keepdim=True)

def report(name, tl_val, nn_val, do_center=True):
    raw = (tl_val - nn_val).abs().max().item()
    line = f"{name:24} raw max|TL-nn| = {raw:8.4f}"
    if do_center:
        cen = (center(tl_val) - center(nn_val)).abs().max().item()
        line += f"   after centering both = {cen:.2e}"
    print(line)

report("resid_pre",  tl_cache["blocks.5.hook_resid_pre"], resid_pre)
report("attn_out",   tl_cache["blocks.5.hook_attn_out"],  attn_out)
report("z",          tl_cache["blocks.5.attn.hook_z"],
       z_in.view(B, S, H, d), do_center=False)
print("   ^ z differs by the value-bias fold; this is expected, not a bug.")

resid_pre                raw max|TL-nn| =   4.5219   after centering both = 1.22e-04
attn_out                 raw max|TL-nn| =   0.0047   after centering both = 7.09e-06
z                        raw max|TL-nn| =   0.5877
   ^ z differs by the value-bias fold; this is expected, not a bug.

The residual stream agrees to ~1e-4 once both sides are centered — the raw difference of 4.5 is entirely the mean that center_writing_weights removed. z does not agree at all, because the value bias has been pushed forward into b_O.

That single distinction explains most "my port is broken" reports.

!!! danger "It bites hardest on weight-space analyses"

Reading activations, you at least notice a discrepancy. Reading weights, a
literal port can look completely reasonable and be wrong by orders of
magnitude — because every published TransformerLens weight statistic was
computed on **folded and centered** weights, while nnsight hands you the raw
checkpoint.

A measured example: a published one-liner for detecting under-trained
("glitch") tokens from the unembedding flags **63** tokens in GPT-2 when run
in TransformerLens. Ported literally onto `model.lm_head.weight`, it flags
**36,604** — most of the vocabulary. Fold `ln_f` into the unembedding and
center it, and the published number comes back.

So before porting anything that reads `W_U`, `W_E`, `W_Q/K/V/O` or a
LayerNorm gain, ask what processing produced the numbers you are comparing
against.

2. Running and caching¶

TransformerLens	nnsight
`logits = model(tokens)`	`with model.trace(tokens): logits = model.lm_head.output.save()`
`logits, cache = model.run_with_cache(t)`	`with model.trace(t) as tracer: cache = tracer.cache(...)`
`run_with_cache(..., names_filter=fn)`	`tracer.cache(modules=[paths or Envoys])` — a list, not a predicate
`run_with_cache(..., stop_at_layer=L)`	`tracer.stop()` after reading layer `L`
`cache["resid_post", 5]`	`cache["model.transformer.h.5"].output`
`HookPoint`, `hook.name`, `hook.layer()`	gone — you address the module directly

One difference worth knowing: nnsight's cache moves values to CPU by default. Pass device=None to keep them where they were computed.

In [4]:

Copied!

with nn.trace(toks) as tracer:
    cache = tracer.cache(modules=[nn.transformer.h[i] for i in range(3)], device=None)

print("cached paths:", list(cache.keys())[:3])
print("layer 0 output:", tuple(cache["model.transformer.h.0"].output.shape))
with nn.trace(toks) as tracer:
    cache = tracer.cache(modules=[nn.transformer.h[i] for i in range(3)], device=None)

print("cached paths:", list(cache.keys())[:3])
print("layer 0 output:", tuple(cache["model.transformer.h.0"].output.shape))

cached paths: ['model.transformer.h.0', 'model.transformer.h.1', 'model.transformer.h.2']
layer 0 output: (1, 15, 768)

3. Hooks become assignments¶

This is the part people enjoy. TransformerLens asks you to write a hook function, name the hook point as a string, and register it. In nnsight you assign to the value.

# TransformerLens: write a hook fn, name the hook point as a string, register it
def patch_resid(resid, hook, src):
    resid[:, -1, :] = src[:, -1, :]
    return resid

logits = model.run_with_hooks(toks, fwd_hooks=[
    ("blocks.8.hook_resid_post", partial(patch_resid, src=corrupt))])

# nnsight: assign to the value
with model.trace(toks):
    model.transformer.h[8].output[:, -1, :] = corrupt[:, -1, :]
    logits = model.lm_head.output.save()

No hook function, no partial, no hook.name, and nothing to reset afterwards — a trace is scoped.

Below we patch the residual stream from a corrupted run (Mary swapped for John) in both libraries and compare the effect on the logit difference. Patching is a behavioural measurement, so the two should agree despite every convention difference in §1.

In [5]:

Copied!





from functools import partial

MARY = nn.tokenizer.encode(" Mary")[0]
JOHN = nn.tokenizer.encode(" John")[0]
logit_diff = lambda lg: (lg[0, -1, MARY] - lg[0, -1, JOHN]).item()

CORRUPT = "When John and Mary went to the shops, Mary gave the bag to"
c_toks = tl.to_tokens([CORRUPT]).to(DEVICE)
PL = 8   # patch the residual stream after block 8, at the last position

tl_base = logit_diff(tl(toks))
with nn.trace(toks):
    nn_base_o = nn.lm_head.output.save()
nn_base = logit_diff(nn_base_o)

# --- TransformerLens: write a hook function, name the hook point, register it ---
def patch_resid(resid, hook, src):
    resid[:, -1, :] = src[:, -1, :]
    return resid

_, tl_corrupt = tl.run_with_cache(c_toks)
tl_patched = logit_diff(tl.run_with_hooks(toks, fwd_hooks=[
    (f"blocks.{PL}.hook_resid_post",
     partial(patch_resid, src=tl_corrupt[f"blocks.{PL}.hook_resid_post"]))]))

# --- nnsight: assign to the value ---
with nn.trace(c_toks):
    corrupt_resid = nn.transformer.h[PL].output.save()
with nn.trace(toks):
    nn.transformer.h[PL].output[:, -1, :] = corrupt_resid[:, -1, :]
    nn_patched_o = nn.lm_head.output.save()
nn_patched = logit_diff(nn_patched_o)

print(f"baseline logit diff      TL {tl_base:+.4f}   nnsight {nn_base:+.4f}")
print(f"after patching layer {PL}   TL {tl_patched:+.4f}   nnsight {nn_patched:+.4f}")
print(f"patching effect          TL {tl_base - tl_patched:+.4f}   "
      f"nnsight {nn_base - nn_patched:+.4f}")
from functools import partial

MARY = nn.tokenizer.encode(" Mary")[0]
JOHN = nn.tokenizer.encode(" John")[0]
logit_diff = lambda lg: (lg[0, -1, MARY] - lg[0, -1, JOHN]).item()

CORRUPT = "When John and Mary went to the shops, Mary gave the bag to"
c_toks = tl.to_tokens([CORRUPT]).to(DEVICE)
PL = 8   # patch the residual stream after block 8, at the last position

tl_base = logit_diff(tl(toks))
with nn.trace(toks):
    nn_base_o = nn.lm_head.output.save()
nn_base = logit_diff(nn_base_o)

# --- TransformerLens: write a hook function, name the hook point, register it ---
def patch_resid(resid, hook, src):
    resid[:, -1, :] = src[:, -1, :]
    return resid

_, tl_corrupt = tl.run_with_cache(c_toks)
tl_patched = logit_diff(tl.run_with_hooks(toks, fwd_hooks=[
    (f"blocks.{PL}.hook_resid_post",
     partial(patch_resid, src=tl_corrupt[f"blocks.{PL}.hook_resid_post"]))]))

# --- nnsight: assign to the value ---
with nn.trace(c_toks):
    corrupt_resid = nn.transformer.h[PL].output.save()
with nn.trace(toks):
    nn.transformer.h[PL].output[:, -1, :] = corrupt_resid[:, -1, :]
    nn_patched_o = nn.lm_head.output.save()
nn_patched = logit_diff(nn_patched_o)

print(f"baseline logit diff      TL {tl_base:+.4f}   nnsight {nn_base:+.4f}")
print(f"after patching layer {PL}   TL {tl_patched:+.4f}   nnsight {nn_patched:+.4f}")
print(f"patching effect          TL {tl_base - tl_patched:+.4f}   "
      f"nnsight {nn_base - nn_patched:+.4f}")

baseline logit diff      TL +3.3367   nnsight +3.3367
after patching layer 8   TL -2.2522   nnsight -2.2522
patching effect          TL +5.5889   nnsight +5.5889

Both libraries agree on the behavioural quantity even though they disagree on the raw value of z — exactly the split from the table at the top.

4. Hook name → module path¶

Verified at block 5 on GPT-2. B=batch, S=seq, H=12 heads, d=64, D=768.

TL hook name	nnsight	note
`hook_embed`	`transformer.wte.output`
`hook_pos_embed`	`transformer.wpe.output`	may come back `[1,S,D]`
`blocks.L.hook_resid_pre`	`transformer.h[L].input`	centering
`blocks.L.hook_resid_mid`	`transformer.h[L].ln_2.input`	centering
`blocks.L.hook_resid_post`	`transformer.h[L].output`	centering
`blocks.L.ln1.hook_normalized`	`(h[L].ln_1.output − ln_1.bias) / ln_1.weight`	raw `ln_1.output` is off by ~24
`blocks.L.attn.hook_q`	`attn.c_attn.output[..., :768].view(B,S,H,d)`	exact
`blocks.L.attn.hook_k`	`attn.c_attn.output[..., 768:1536].view(B,S,H,d)`	exact
`blocks.L.attn.hook_v`	`attn.c_attn.output[..., 1536:].view(B,S,H,d)`	value-bias fold
`blocks.L.attn.hook_z`	`attn.c_proj.input.view(B,S,H,d)`	value-bias fold
`blocks.L.attn.hook_pattern`	`attn.source.attention_interface_1.output[1]`	needs eager attention
`blocks.L.attn.hook_attn_scores`	`attn.source.attention_interface_1.source.nn_functional_softmax_0.input`	needs eager attention
`blocks.L.hook_attn_out`	`transformer.h[L].attn.output[0]`	tuple
`blocks.L.hook_mlp_out`	`transformer.h[L].mlp.output`	tensor
`blocks.L.attn.hook_result`	`einsum("bshd,hdm->bshm", z, c_proj.weight.view(H,d,D))`

.source is how you reach a value computed inside a module's forward, which TransformerLens can only offer because it reimplemented the architecture. Print model.transformer.h[5].attn.source to see the actual forward with its hook points labelled.

In [6]:

Copied!





# Warm-up (see Gotchas): the first `.source` access on a module, in a trace that
# already read something else, can raise OutOfOrderError. A throwaway trace that
# touches the source on its own resolves it.
with nn.trace(toks):
    _ = nn.transformer.h[L].attn.source.attention_interface_1.output[1].save()

with nn.trace(toks):
    q       = nn.transformer.h[L].attn.c_attn.output[..., :768].save()
    pattern = nn.transformer.h[L].attn.source.attention_interface_1.output[1].save()
    mlp_out = nn.transformer.h[L].mlp.output.save()

print("q       max|TL-nn| =", (tl_cache["blocks.5.attn.hook_q"]
                               - q.view(B, S, H, d)).abs().max().item())
print("pattern max|TL-nn| =", (tl_cache["blocks.5.attn.hook_pattern"]
                               - pattern).abs().max().item())
print("mlp_out max|TL-nn| =", (center(tl_cache["blocks.5.hook_mlp_out"])
                               - center(mlp_out)).abs().max().item(), "(centered)")
# Warm-up (see Gotchas): the first `.source` access on a module, in a trace that
# already read something else, can raise OutOfOrderError. A throwaway trace that
# touches the source on its own resolves it.
with nn.trace(toks):
    _ = nn.transformer.h[L].attn.source.attention_interface_1.output[1].save()

with nn.trace(toks):
    q       = nn.transformer.h[L].attn.c_attn.output[..., :768].save()
    pattern = nn.transformer.h[L].attn.source.attention_interface_1.output[1].save()
    mlp_out = nn.transformer.h[L].mlp.output.save()

print("q       max|TL-nn| =", (tl_cache["blocks.5.attn.hook_q"]
                               - q.view(B, S, H, d)).abs().max().item())
print("pattern max|TL-nn| =", (tl_cache["blocks.5.attn.hook_pattern"]
                               - pattern).abs().max().item())
print("mlp_out max|TL-nn| =", (center(tl_cache["blocks.5.hook_mlp_out"])
                               - center(mlp_out)).abs().max().item(), "(centered)")

q       max|TL-nn| = 3.993511199951172e-06
pattern max|TL-nn| = 1.1920928955078125e-06
mlp_out max|TL-nn| = 4.57763671875e-05 (centered)

5. What you gain: many interventions, one forward pass¶

TransformerLens runs one run_with_hooks per intervention. nnsight lets you put several invokes in a single trace, which become one batched forward — with different interventions on different rows.

For a 12-layer sweep that is roughly a 7x wall-clock win below, with results identical to the one-at-a-time version to 5e-5.

In [7]:

Copied!





# cache the corrupted run once
with nn.trace(c_toks) as tracer:
    c_cache = tracer.cache(modules=[nn.transformer.h[i] for i in range(12)], device=None)
corrupt = [c_cache[f"model.transformer.h.{i}"].output for i in range(12)]

# --- one trace, 12 invokes, one batched forward ---
t0 = time.time()
with nn.trace() as tracer:
    outs = nnsight.save([])          # save the CONTAINER, append raw values into it
    for i in range(12):
        with tracer.invoke(toks):
            nn.transformer.h[i].output[:, -1, :] = corrupt[i][:, -1, :]
            outs.append(nn.lm_head.output)
batched_s = time.time() - t0
batched = [logit_diff(o) for o in outs]

# --- the TransformerLens shape: one forward per layer ---
t0 = time.time()
loop = []
for i in range(12):
    with nn.trace(toks):
        nn.transformer.h[i].output[:, -1, :] = corrupt[i][:, -1, :]
        o = nn.lm_head.output.save()
    loop.append(logit_diff(o))
loop_s = time.time() - t0

print(f"12 invokes in one trace : {batched_s*1000:6.1f} ms")
print(f"12 separate traces      : {loop_s*1000:6.1f} ms")
print(f"max difference in result: {max(abs(a-b) for a, b in zip(batched, loop)):.2e}")
print("\nlayer :", " ".join(f"{i:6d}" for i in range(12)))
print("effect:", " ".join(f"{v:6.2f}" for v in batched))
# cache the corrupted run once
with nn.trace(c_toks) as tracer:
    c_cache = tracer.cache(modules=[nn.transformer.h[i] for i in range(12)], device=None)
corrupt = [c_cache[f"model.transformer.h.{i}"].output for i in range(12)]

# --- one trace, 12 invokes, one batched forward ---
t0 = time.time()
with nn.trace() as tracer:
    outs = nnsight.save([])          # save the CONTAINER, append raw values into it
    for i in range(12):
        with tracer.invoke(toks):
            nn.transformer.h[i].output[:, -1, :] = corrupt[i][:, -1, :]
            outs.append(nn.lm_head.output)
batched_s = time.time() - t0
batched = [logit_diff(o) for o in outs]

# --- the TransformerLens shape: one forward per layer ---
t0 = time.time()
loop = []
for i in range(12):
    with nn.trace(toks):
        nn.transformer.h[i].output[:, -1, :] = corrupt[i][:, -1, :]
        o = nn.lm_head.output.save()
    loop.append(logit_diff(o))
loop_s = time.time() - t0

print(f"12 invokes in one trace : {batched_s*1000:6.1f} ms")
print(f"12 separate traces      : {loop_s*1000:6.1f} ms")
print(f"max difference in result: {max(abs(a-b) for a, b in zip(batched, loop)):.2e}")
print("\nlayer :", " ".join(f"{i:6d}" for i in range(12)))
print("effect:", " ".join(f"{v:6.2f}" for v in batched))

12 invokes in one trace :   12.8 ms
12 separate traces      :   87.9 ms
max difference in result: 5.34e-05

layer :      0      1      2      3      4      5      6      7      8      9     10     11
effect:   3.34   3.35   3.34   3.34   3.33   3.33   3.24   1.06  -2.25  -2.76  -2.71  -3.20

6. Gotchas worth pinning to the wall¶

.output is whatever the module really returns. On transformers 5.x a decoder block returns a bare tensor; an attention submodule returns a tuple. Indexing a tensor with [0] does not raise — it silently gives you the first batch element. Check with with model.trace(x): print(type(module.output)).
Tokenization. to_tokens prepends BOS; tokenizer(...) does not.
Attention probabilities need attn_implementation="eager". Under the default SDPA they come back None.
tracer.cache(...) moves values to CPU unless you pass device=None.
No weight processing. There is no fold_ln / center_writing_weights / fold_value_biases equivalent, by design — nnsight shows you the checkpoint.
Nested .source needs one warm-up. The first access to a nested source op on a given module, in a trace that already read something else, can raise OutOfOrderError; the identical block succeeds on the next attempt. Touch it in a throwaway trace first.

What has no nnsight equivalent¶

These are the rows where a port needs you to write code, not translate it:

ActivationCache decomposition algebra — accumulated_resid, decompose_resid, stack_head_results, apply_ln_to_stack
the weight-processing flags (§1)
utils.test_prompt, model.get_token_position
a normalised cfg — you read HuggingFace config names per family

If your work leans on those, nnterp provides a standardized layer over nnsight with architecture-agnostic accessors.