Info

Last Execution: 2026-09-02

Package	Version
nnsight	0.8.0
Python	3.12.13
torch	2.13.0+cu126
transformers	5.15.0.dev0

Module Access¶

nnsight wraps every model in an Envoy tree that mirrors the underlying torch.nn.Module hierarchy. Every submodule is reachable by the same attribute path you would use in plain PyTorch (model.transformer.h[0].mlp), and each one exposes its live values during a forward pass through .output, .input, and .inputs.

This page covers how to find your way around that tree: printing it, navigating by attribute and index, giving modules portable aliases with rename=, calling a module directly as a function, and reaching operations inside a forward with .source.

In [1]:

from nnsight.modeling.transformers import TransformersModel
import torch

model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)

from nnsight.modeling.transformers import TransformersModel import torch model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)

Printing the Envoy Tree¶

Just like a normal PyTorch nn.Module, you print an Envoy to see its structure: print(model) renders the whole module tree. This is your map: every name and index shown here is a valid attribute path for accessing activations.

In [2]:

print(model)

print(model)

GPT2LMHeadModel(
  (transformer): GPT2Model(
    (wte): Embedding(50257, 768)
    (wpe): Embedding(1024, 768)
    (drop): Dropout(p=0.1, inplace=False)
    (h): ModuleList(
      (0-11): 12 x GPT2Block(
        (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
        (attn): GPT2Attention(
          (c_attn): Conv1D()
          (c_proj): Conv1D()
          (attn_dropout): Dropout(p=0.1, inplace=False)
          (resid_dropout): Dropout(p=0.1, inplace=False)
        )
        (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
        (mlp): GPT2MLP(
          (c_fc): Conv1D()
          (c_proj): Conv1D()
          (act): NewGELUActivation()
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
    )
    (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
  )
  (lm_head): Linear(in_features=768, out_features=50257, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
)

Navigating by Attribute and Index¶

Read any name off the printed tree to reach that Envoy. ModuleLists (like the stack of transformer blocks) are indexed, and negative indices work as usual.

In [3]:

# A single block, its attention submodule, and its MLP.
block = model.transformer.h[0]
print(block.attn)
print(block.mlp)

# A single block, its attention submodule, and its MLP. block = model.transformer.h[0] print(block.attn) print(block.mlp)

GPT2Attention(
  (c_attn): Conv1D()
  (c_proj): Conv1D()
  (attn_dropout): Dropout(p=0.1, inplace=False)
  (resid_dropout): Dropout(p=0.1, inplace=False)
)
GPT2MLP(
  (c_fc): Conv1D()
  (c_proj): Conv1D()
  (act): NewGELUActivation()
  (dropout): Dropout(p=0.1, inplace=False)
)

Renaming Modules¶

Different architectures name the same role differently (transformer.h vs model.layers vs gpt_neox.layers). Pass rename={...} at construction to install aliases so your intervention code is portable across model families. An alias points at the same Envoy object, so the original path keeps working too.

In [4]:

renamed = TransformersModel(
    "openai-community/gpt2",
    device_map="auto",
    dispatch=True,
    rename={
        "transformer.h": "layers",  # mount a subtree at a shorter path
        "mlp": "ffn",               # rename every block's MLP
    },
)

# Alias and original path resolve to the exact same Envoy.
print(renamed.layers[0].ffn is renamed.transformer.h[0].mlp)

with renamed.trace("Hello"):
    a = renamed.layers[0].ffn.output.save()       # via aliases
    b = renamed.transformer.h[0].mlp.output.save()  # original still works

print(torch.equal(a, b))

renamed = TransformersModel( "openai-community/gpt2", device_map="auto", dispatch=True, rename={ "transformer.h": "layers", # mount a subtree at a shorter path "mlp": "ffn", # rename every block's MLP }, ) # Alias and original path resolve to the exact same Envoy. print(renamed.layers[0].ffn is renamed.transformer.h[0].mlp) with renamed.trace("Hello"): a = renamed.layers[0].ffn.output.save() # via aliases b = renamed.transformer.h[0].mlp.output.save() # original still works print(torch.equal(a, b))

True

True

Calling a Module as a Function¶

Inside a trace you can call any module as a function on values you already have. The call runs with this trace stood down, so it does not consume the module's place in the forward pass and is not bound by execution order. That is what makes it useful for applying a module out of place.

In [5]:

with model.trace("The Eiffel Tower is in the city of"):

    hs = model.transformer.h[-1].output

    # Apply the final layer norm then lm_head to decode the hidden states.
    logits = model.lm_head(model.transformer.ln_f(hs))
    token = logits[0, -1].argmax(dim=-1).save()

print(f"Decoded: {model.tokenizer.decode(token)}")

with model.trace("The Eiffel Tower is in the city of"): hs = model.transformer.h[-1].output # Apply the final layer norm then lm_head to decode the hidden states. logits = model.lm_head(model.transformer.ln_f(hs)) token = logits[0, -1].argmax(dim=-1).save() print(f"Decoded: {model.tokenizer.decode(token)}")

Decoded:  Paris

Function calls vs. attribute access

When you call a module (e.g. model.lm_head(x)), nnsight stands this trace down for the duration: the call serves no value and spends no occurrence, so .output still refers to the module's real place in the forward pass. The module's own PyTorch hooks are untouched. hook= names whether nnsight watches the call, not whether the module runs its hooks, so a forward_hook you registered yourself fires for this call as well as for the real forward. Pass hook=True when you want nnsight to watch it too, which is what makes an attached module's submodules observable at .submodule.output.

Logit Lens¶

A classic application: decode every layer's hidden states through the final layer norm and lm_head to see what the model "thinks" at each layer.

In [6]:

with model.trace("The Eiffel Tower is in the city of"):

    predictions = list().save()

    for layer in model.transformer.h:
        hs = layer.output
        logits = model.lm_head(model.transformer.ln_f(hs))
        predictions.append(logits[0, -1].argmax(dim=-1))

for i, tok in enumerate(predictions):
    print(f"Layer {i:2d}: {model.tokenizer.decode(tok)}")

with model.trace("The Eiffel Tower is in the city of"): predictions = list().save() for layer in model.transformer.h: hs = layer.output logits = model.lm_head(model.transformer.ln_f(hs)) predictions.append(logits[0, -1].argmax(dim=-1)) for i, tok in enumerate(predictions): print(f"Layer {i:2d}: {model.tokenizer.decode(tok)}")

Layer  0:  the
Layer  1:  the
Layer  2:  the
Layer  3:  the
Layer  4:  the
Layer  5:  the
Layer  6:  East
Layer  7:  Ing
Layer  8:  Rome
Layer  9:  London
Layer 10:  Paris
Layer 11:  Paris

Notice how the prediction converges toward "Paris" in the later layers — this is the logit lens in action. Each layer's intermediate representation is projected into vocabulary space using the same final modules.

Combining Module Calls with Torch Operations¶

Everything inside a trace is real PyTorch, so you can freely mix module calls with standard torch operations.

In [7]:

with model.trace("The Eiffel Tower is in the city of"):

    hs = model.transformer.h[-1].output
    normed = model.transformer.ln_f(hs)

    # Cosine similarity between the last token and every earlier token.
    similarity = torch.cosine_similarity(
        normed[:, -1:, :], normed[:, :-1, :], dim=-1
    ).save()

print(f"Cosine similarity with last token: {similarity[0]}")

with model.trace("The Eiffel Tower is in the city of"): hs = model.transformer.h[-1].output normed = model.transformer.ln_f(hs) # Cosine similarity between the last token and every earlier token. similarity = torch.cosine_similarity( normed[:, -1:, :], normed[:, :-1, :], dim=-1 ).save() print(f"Cosine similarity with last token: {similarity[0]}")

Cosine similarity with last token: tensor([0.9684, 0.9578, 0.9684, 0.9637, 0.9910, 0.9959, 0.9829, 0.9934, 0.9821],
       device='cuda:0', grad_fn=<SelectBackward0>)

To reach operations inside a module's forward — the values that live between two operations, with no submodule to attach to — use module.source. See Intermediate Operations for the full story.

Other Ways to Reach Modules¶

Beyond attribute-and-index navigation, an Envoy exposes a handful of helpers for reading the underlying module directly, resolving paths built at runtime, and walking the tree programmatically.

Reading the Underlying Module¶

An Envoy forwards any attribute it doesn't recognize to the nn.Module it wraps. So you can read the module's own attributes and parameters straight off the Envoy — model.transformer.wte.weight hands back the embedding weight tensor. Accessing a submodule this way auto-wraps it as an Envoy, so it keeps behaving like the rest of the tree.

In [8]:

# The wrapped module's parameter, read directly through the Envoy.
wte_weight = model.transformer.wte.weight
print(type(wte_weight).__name__, tuple(wte_weight.shape))

# The wrapped module's parameter, read directly through the Envoy. wte_weight = model.transformer.wte.weight print(type(wte_weight).__name__, tuple(wte_weight.shape))

Parameter (50257, 768)

Name Collisions¶

nnsight puts .output, .input, and .inputs on every Envoy. But some architectures name a submodule the same as one of these — BERT's encoder layers each contain a submodule literally named output. When that happens the submodule wins the name and nnsight's attribute moves aside to nns_<name> (a warning is emitted when the tree is built). So on a BERT layer, .output is the submodule and the layer's intervention output lives at .nns_output.

In [9]:

import warnings

with warnings.catch_warnings():
    warnings.simplefilter("ignore")  # silence the expected shadow warning
    bert = TransformersModel(
        "google-bert/bert-base-uncased", device_map="auto", dispatch=True
    )

layer = bert.bert.encoder.layer[0]
# `.output` now wraps the BertOutput submodule that shadowed Envoy's `.output`;
# nnsight's intervention output has moved aside to `.nns_output`.
print("layer.output wraps: ", type(layer.output._module).__name__)
print("layer.nns_output is:", type(type(layer).nns_output).__name__)

import warnings with warnings.catch_warnings(): warnings.simplefilter("ignore") # silence the expected shadow warning bert = TransformersModel( "google-bert/bert-base-uncased", device_map="auto", dispatch=True ) layer = bert.bert.encoder.layer[0] # `.output` now wraps the BertOutput submodule that shadowed Envoy's `.output`; # nnsight's intervention output has moved aside to `.nns_output`. print("layer.output wraps: ", type(layer.output._module).__name__) print("layer.nns_output is:", type(type(layer).nns_output).__name__)

[transformers] BertForMaskedLM LOAD REPORT from: google-bert/bert-base-uncased
Key                         | Status     |  | 
----------------------------+------------+--+-
cls.seq_relationship.bias   | UNEXPECTED |  | 
bert.pooler.dense.weight    | UNEXPECTED |  | 
cls.seq_relationship.weight | UNEXPECTED |  | 
bert.pooler.dense.bias      | UNEXPECTED |  | 

Notes:
- UNEXPECTED:	can be ignored when loading from different task/architecture; not ok if you expect identical arch.

layer.output wraps:  BertOutput
layer.nns_output is: eproperty

To read or patch that layer's intervention output you reach for .nns_output, because .output now points at the submodule. (This BERT is a fill-mask model, so the prompt needs a [MASK] token.)

In [10]:

with bert.trace("The Eiffel Tower is in the city of [MASK]."):
    hs = layer.nns_output.save()

print(f"Layer 0 intervention output: {tuple(hs.shape)}")

with bert.trace("The Eiffel Tower is in the city of [MASK]."): hs = layer.nns_output.save() print(f"Layer 0 intervention output: {tuple(hs.shape)}")

Layer 0 intervention output: (1, 14, 768)

Where a Module Lives: .device and .devices¶

.device is the device of the module's first parameter (or None if it has none); .devices is the set of devices its parameters are spread across — handy for a model sharded with device_map="auto".

In [11]:

print("device: ", model.device)
print("devices:", model.devices)

print("device: ", model.device) print("devices:", model.devices)

device:  cuda:0
devices: {device(type='cuda', index=0)}

Resolving a Path at Runtime: .get()¶

When the path you want is built at runtime rather than written out, .get() resolves a dotted string to the same Envoy attribute access would. Outside a trace it returns the descendant Envoy; inside one, a trailing .output/.input resolves through to the live value.

In [12]:

# Outside a trace: resolves to the child Envoy.
print(model.get("transformer.h.0.mlp"))

with model.trace("The Eiffel Tower is in the city of"):
    # Inside a trace: a trailing `.output` resolves to the live value.
    mlp_out = model.get("transformer.h.0.mlp.output").save()

print(f"MLP output: {tuple(mlp_out.shape)}")

# Outside a trace: resolves to the child Envoy. print(model.get("transformer.h.0.mlp")) with model.trace("The Eiffel Tower is in the city of"): # Inside a trace: a trailing `.output` resolves to the live value. mlp_out = model.get("transformer.h.0.mlp.output").save() print(f"MLP output: {tuple(mlp_out.shape)}")

GPT2MLP(
  (c_fc): Conv1D()
  (c_proj): Conv1D()
  (act): NewGELUActivation()
  (dropout): Dropout(p=0.1, inplace=False)
)
MLP output: (1, 10, 768)

Walking the Tree: .modules() and .named_modules()¶

.modules() flattens the whole subtree (children first, then self) into a flat list of Envoys; .named_modules() gives (path, Envoy) tuples instead. Both take an optional include_fn predicate to keep only the Envoys you want.

In [13]:

print(f"Total modules: {len(model.modules())}")

# First few (path, Envoy) pairs.
for path, envoy in model.named_modules()[:5]:
    print(path)

# Filter with a predicate: just the MLP modules, one per block.
mlps = model.named_modules(include_fn=lambda e: e.path.endswith(".mlp"))
print(f"\nMLP modules: {len(mlps)} (e.g. {mlps[0][0]})")

print(f"Total modules: {len(model.modules())}") # First few (path, Envoy) pairs. for path, envoy in model.named_modules()[:5]: print(path) # Filter with a predicate: just the MLP modules, one per block. mlps = model.named_modules(include_fn=lambda e: e.path.endswith(".mlp")) print(f"\nMLP modules: {len(mlps)} (e.g. {mlps[0][0]})")

Total modules: 166
model.transformer.wte
model.transformer.wpe
model.transformer.drop
model.transformer.h.0.ln_1
model.transformer.h.0.attn.c_attn

MLP modules: 12 (e.g. model.transformer.h.0.mlp)

Gotchas¶

A couple of rules govern how you access the tree.

Access modules in forward-pass order within an invoke. Requesting a later module and then an earlier one deadlocks and raises OutOfOrderError. To read modules out of order, use separate invokes.

In [14]:

from nnsight.intervention.interleaver import OutOfOrderError

try:
    with model.trace("Hello world"):
        later = model.transformer.h[5].output.save()
        earlier = model.transformer.h[2].output.save()  # runs before h[5]
except OutOfOrderError as e:
    print(f"OutOfOrderError: {e}")

from nnsight.intervention.interleaver import OutOfOrderError try: with model.trace("Hello world"): later = model.transformer.h[5].output.save() earlier = model.transformer.h[2].output.save() # runs before h[5] except OutOfOrderError as e: print(f"OutOfOrderError: {e}")

OutOfOrderError: 'model.transformer.h.2.output.i0' was requested but the model already ran past it

.save() only makes sense inside a trace. Called outside a tracing context it raises, because there is nothing for it to return from. Inside a trace it marks a value to survive past the with block, and the value comes back under the variable name you assigned it to.

In [15]:

try:
    torch.zeros(3).save()
except ValueError as e:
    print(f"ValueError: {e}")

try: torch.zeros(3).save() except ValueError as e: print(f"ValueError: {e}")

ValueError: save() was called outside a trace. `.save()` / nnsight.save(x) marks a value to return from the enclosing `with model.trace(...):` block, so it only works inside one — move the save into the trace block.