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

Getting Activations¶

This tutorial assumes you've already been through the main Walkthrough.

Reading intermediate values from a model's forward pass is one of the most fundamental operations in nnsight. This page covers how to access a module's output, its inputs, and how to persist those values so you can use them after the trace ends.

Setup¶

We load GPT-2 with TransformersModel.

More on models

For details on TransformersModel, device_map, dispatching, and loading options, see the transformers model guide and the Loading a Model tutorial.

In [1]:

from nnsight.modeling.transformers import TransformersModel

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

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

Getting a Layer Output¶

Open a tracing context with model.trace(...), then read a module's forward-pass return value with .output.

What .output hands you inside the block is the real tensor: .shape, .mean() and .item() all work on it right there. What it does not do is survive the block. The body runs in a worker alongside the forward pass, and only the values you mark come back. .save() is that mark, and a saved value returns under the name you bound it to. This is your first look at .save(); it shows up throughout these tutorials.

In [2]:

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

    hidden_states = model.transformer.h[-1].output.save()

print(hidden_states.shape)

with model.trace("The Eiffel Tower is in the city of"): hidden_states = model.transformer.h[-1].output.save() print(hidden_states.shape)

torch.Size([1, 10, 768])

Always .save() and bind to a variable

A traced with block does not run in place — its body runs in a worker that only returns values you explicitly mark. .save() marks a value to survive the block, and it comes back by the variable name you bind it to, so write hidden_states = ....save(). A bare model.transformer.h[-1].output.save() on its own line is marked but has no name to return under, so it is silently lost.

Getting a Layer Input¶

Use .input to read the first positional argument passed into a module.

In [3]:

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

    layer_input = model.transformer.h[0].input.save()

print(layer_input.shape)

with model.trace("The Eiffel Tower is in the city of"): layer_input = model.transformer.h[0].input.save() print(layer_input.shape)

torch.Size([1, 10, 768])

.input vs .inputs¶

.input is a convenience returning just the first positional argument. .inputs returns the full call signature as an (args, kwargs) tuple — every positional argument and every keyword argument the module was called with.

In [4]:

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

    detailed_inputs = model.transformer.h[0].inputs.save()

args, kwargs = detailed_inputs
print(f"Positional args: {len(args)}")
print(f"Keyword args: {list(kwargs.keys())}")

with model.trace("The Eiffel Tower is in the city of"): detailed_inputs = model.transformer.h[0].inputs.save() args, kwargs = detailed_inputs print(f"Positional args: {len(args)}") print(f"Keyword args: {list(kwargs.keys())}")

Positional args: 4
Keyword args: ['encoder_attention_mask', 'use_cache', 'position_ids']

Access modules in forward-pass order

Within a single invoke, request modules in the order they run. Reading a later module and then an earlier one deadlocks and raises OutOfOrderError once the forward pass finishes past the earlier module. To read modules out of order, use separate invokes (see the batching guide).

nnsight.save() vs .save()¶

There are two equivalent ways to save a value. The preferred form is the function nnsight.save(...), which works on any object:

In [5]:

import nnsight

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

    hidden_states = nnsight.save(model.transformer.h[-1].output)

print(hidden_states.shape)

import nnsight with model.trace("The Eiffel Tower is in the city of"): hidden_states = nnsight.save(model.transformer.h[-1].output) print(hidden_states.shape)

torch.Size([1, 10, 768])

nnsight.save() vs obj.save()

obj.save() relies on a C extension that mounts a .save() method onto every Python object at import time. It works in most cases, but nnsight.save() is safer: it is unaffected if a class defines its own .save() method (which would shadow nnsight's version). Prefer nnsight.save() for plain Python values like ints, lists, and dicts.

Getting Multiple Layer Outputs¶

Collect hidden states from every layer in a single forward pass. Save the container and put the raw per-layer outputs into it — do not .save() the individual elements.

In [6]:

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

    hidden_states_per_layer = nnsight.save(
        [layer.output for layer in model.transformer.h]
    )

for i, hs in enumerate(hidden_states_per_layer):
    print(f"Layer {i}: {hs.shape}")

with model.trace("The Eiffel Tower is in the city of"): hidden_states_per_layer = nnsight.save( [layer.output for layer in model.transformer.h] ) for i, hs in enumerate(hidden_states_per_layer): print(f"Layer {i}: {hs.shape}")

Layer 0: torch.Size([1, 10, 768])
Layer 1: torch.Size([1, 10, 768])
Layer 2: torch.Size([1, 10, 768])
Layer 3: torch.Size([1, 10, 768])
Layer 4: torch.Size([1, 10, 768])
Layer 5: torch.Size([1, 10, 768])
Layer 6: torch.Size([1, 10, 768])
Layer 7: torch.Size([1, 10, 768])
Layer 8: torch.Size([1, 10, 768])
Layer 9: torch.Size([1, 10, 768])
Layer 10: torch.Size([1, 10, 768])
Layer 11: torch.Size([1, 10, 768])

Saving collections

Save the list itself (here via nnsight.save([...])); a saved container comes back with its contents. If you instead .save() each element the marks have no name to return under, and if you leave the list unsaved it never comes back at all.

Using Saved Outputs¶

Saved activations are real tensors, so you can run any PyTorch operation on them after the tracing context exits.

In [7]:

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

    logits = model.lm_head.output.save()

predicted_token = logits[0, -1].argmax(dim=-1)
print(f"Predicted next token: {model.tokenizer.decode(predicted_token)}")

with model.trace("The Eiffel Tower is in the city of"): logits = model.lm_head.output.save() predicted_token = logits[0, -1].argmax(dim=-1) print(f"Predicted next token: {model.tokenizer.decode(predicted_token)}")

Predicted next token:  Paris