Info

Last Execution: 2026-02-17

Package	Version
nnsight	0.5.15
Python	3.12.3
torch	2.10.0+cu128
transformers	5.2.0

Getting Activations¶

Accessing intermediate values from a model's forward pass is one of the most fundamental operations in nnsight. This page covers how to read hidden states, inputs, and outputs at any layer.

Setup¶

In [1]:

from nnsight import LanguageModel

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

from nnsight import LanguageModel model = LanguageModel("openai-community/gpt2", device_map="auto", dispatch=True)

Getting a Layer Output¶

Access any module's output by calling .output on it inside a tracing context. Use .save() to persist the value after the context exits.

In [2]:

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

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

print(hidden_states.shape)

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

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

Always use .save()

Without .save(), values are automatically cleaned up when the tracing context exits. Always call .save() on any value you need afterward.

Getting a Layer Input¶

Use .input to access the first positional argument passed to 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 / .output vs .inputs / .outputs¶

.input returns just the first positional argument. .inputs returns the full signature as a tuple of (args_tuple, kwargs_dict):

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: 5
Keyword args: ['encoder_attention_mask', 'use_cache', 'output_attentions', 'position_ids']

nnsight.save() vs .save()¶

There are two ways to save values. The preferred method is 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[0])

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[0]) print(hidden_states.shape)

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

nnsight.save() vs obj.save()

obj.save() is a convenience method that relies on a C extension (pymount) to inject .save() onto every Python object at runtime. It works in most cases, but nnsight.save() is safer because it works even on objects that define their own .save() method (which would shadow nnsight's version). Use nnsight.save() for non-tensor values like shape integers, lists, or dicts.

Getting Multiple Layer Outputs¶

Loop over layers to collect hidden states from every layer in a single forward pass.

In [6]:

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

    hidden_states_per_layer = list().save()

    for layer in model.transformer.h:
        hidden_states_per_layer.append(layer.output[0])

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 = list().save() for layer in model.transformer.h: hidden_states_per_layer.append(layer.output[0]) 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

Call .save() on the list itself so it persists after the context exits. The individual tensors appended inside will be retained as part of the saved list.

Using Saved Outputs¶

Saved activations are real tensors. You can use them with standard PyTorch operations after the tracing context.

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