Nnsight Basics

NNsight Basics

NNsight provides full access to neural network internals during forward passes. This skill covers the fundamental patterns for tracing, saving, and modifying activations.

Installation

pip install nnsight
pip install torch transformers

Loading Models

For language models (recommended):

from nnsight import LanguageModel

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

For any PyTorch model:

from nnsight import NNsight

model = NNsight(your_pytorch_model)

The Tracing Context

All interventions happen inside model.trace(). Your code interleaves with the forward pass.

with model.trace("The Eiffel Tower is in"):
    # Access any layer's output
    hidden_states = model.transformer.h[5].output[0].save()

    # Access final logits
    logits = model.lm_head.output.save()

# Values available after tracing
print(hidden_states.value.shape)

Accessing Module Values

with model.trace(prompt):
    layer_output = model.transformer.h[0].output[0].save()
    layer_input = model.transformer.h[0].input[0].save()
    attn_output = model.transformer.h[0].attn.output[0].save()

Saving Values

Values only persist if explicitly saved:

# Call .save() on the proxy to persist it outside the trace
output = model.transformer.h[0].output[0].save()

After exiting the trace context, access the actual tensor via .value:

print(output.value.shape)  # Access the saved tensor

Caching Activations in One Pass

Every with model.trace(...): is a full forward pass. If you want activations from many modules of the same input, gather them inside a single trace by looping over the modules — don't open one trace per module.

# One forward pass, collects every layer's output
hidden_states = list()
with model.trace(prompt):
    for layer in model.transformer.h:
        hidden_states.append(layer.output[0].save())

print(hidden_states[0].value.shape)

The anti-pattern is opening N traces to fetch N layers' outputs from one input — that's N forward passes for no reason.

For the catch-all "give me everything" case, tracer.cache() captures all module outputs without listing them by name:

with model.trace(prompt) as tracer:
    cache = tracer.cache()

The general rule: structure traces by what input you're running, not by what activations you're gathering.

Basic Interventions

Zero out activations:

with model.trace(prompt):
    model.transformer.h[5].output[0][:, :, :] = 0
    output = model.lm_head.output.save()

Add noise:

with model.trace(prompt):
    noise = 0.1 * torch.randn_like(model.transformer.h[5].output[0])
    model.transformer.h[5].output[0] = model.transformer.h[5].output[0] + noise

Replace with custom values:

with model.trace(prompt):
    model.transformer.h[5].output[0][:, -1, :] = my_custom_vector

Batched Processing with Invokers

tracer.invoke() runs multiple inputs through one forward pass while keeping their interventions isolated. Indices inside an invoke block refer to that invoke's slice of the batch — write [:, -1, :] as if it were the only input and nnsight handles the underlying batch arithmetic.

Distinct prompts in one pass:

with model.trace() as tracer:
    with tracer.invoke("First prompt"):
        first_output = model.lm_head.output.save()

    with tracer.invoke("Second prompt"):
        second_output = model.lm_head.output.save()

print(first_output.value.shape, second_output.value.shape)

Sweep pattern. When you want the same intervention applied with one varying parameter (patching each head one at a time, ablating each position, sweeping a steering coefficient), put the loop inside the trace and create one invoke per variant. Don't pass the input to model.trace() — pass it to each invoke:

results = list().save()
with model.trace() as tracer:                        # no input here
    for head_idx in range(n_heads):
        with tracer.invoke(prompt):                  # input passed per-invoke
            start, end = head_idx * head_dim, (head_idx + 1) * head_dim
            model.transformer.h[layer].attn.c_proj.input[0][0][:, :, start:end] = \
                clean_attn[:, :, start:end]
            results.append(model.lm_head.output[:, -1].save())

This collapses N forward passes into one. The effective batch size grows with the number of invokes, so chunk the loop into groups if you OOM.

Module Skipping

module.skip(value) bypasses a module's forward pass entirely — value becomes the output that propagates to the next module. Use it to substitute cached activations or short-circuit computation you've already done.

# Pre-cache layer outputs from one input
clean_acts = list()
with model.trace(clean_prompt):
    for layer in model.model.layers:
        clean_acts.append(layer.output)

# On a different input, skip layers 0..L-1 by replacing their outputs with the cache,
# then intervene at layer L with everything before it bypassed
with model.trace(other_prompt):
    for i in range(L):
        model.model.layers[i].skip(clean_acts[i])
    # ...patch at layer L...

The savings scale with how much computation you skip — depth × sequence length. On small models with short sequences the bookkeeping can outweigh the savings; on deep models it's substantial.

Multi-invoke constraint. If you call .skip() on a module in any invoke of a trace, you must call it in every invoke of that trace. Invokes share one underlying forward pass, so the skip has to apply uniformly:

with model.trace() as tracer:
    with tracer.invoke(prompt_a):
        model.model.layers[3].skip(value_a)  # required
    with tracer.invoke(prompt_b):
        model.model.layers[3].skip(value_b)  # also required, even if value differs

Cross-Prompt Interventions (Barriers)

To use values from one invoke in another invoke, use barrier() for synchronization:

with model.trace() as tracer:
    # Create barrier for 2 participating invokes
    barrier = tracer.barrier(2)

    with tracer.invoke("The Eiffel Tower is in"):
        embeddings = model.transformer.wte.output
        barrier()  # Signal: value captured

    with tracer.invoke("_ _ _ _ _ _ _"):
        barrier()  # Wait for value from previous invoke
        model.transformer.wte.output = embeddings  # Use captured value
        output = model.lm_head.output.save()

Note: Barriers are only needed when setting values across invokes. Reading values independently in each invoke doesn't require barriers.

Multi-Token Generation

For autoregressive generation with interventions:

with model.generate("Hello", max_new_tokens=5) as tracer:
    # Get the full generated output
    output = model.generator.output.save()

# Decode the generated tokens
print(model.tokenizer.decode(output[0]))

Apply interventions to all generation steps:

with model.generate("Hello", max_new_tokens=5) as tracer:
    hidden_states = list().save()

    with tracer.all():  # Apply to ALL generation iterations
        # Intervention applied at each step
        model.transformer.h[5].output[0][:, -1, :] *= 0.9
        hidden_states.append(model.transformer.h[-1].output)

Apply interventions to specific iterations:

with model.generate("Hello", max_new_tokens=10) as tracer:
    with tracer.iter[2:5]:  # Only iterations 2, 3, 4
        model.transformer.h[5].output[0][:] = 0

Gradient Access

Access gradients using the .backward() context manager. Gradients must be accessed in reverse order (following backpropagation order):

with model.trace(prompt):
    # Register intermediate values in forward order
    hidden = model.transformer.h[-1].output[0]
    hidden.requires_grad = True  # Enable gradient flow

    logits = model.lm_head.output

    # Enter backward context
    with logits.sum().backward():
        # Access gradients in REVERSE order
        logits_grad = logits.grad.save()
        hidden_grad = hidden.grad.save()

Alternative with retain_grad() (simpler for just accessing gradients):

with model.trace(prompt):
    hidden = model.transformer.h[-1].output[0].save()
    hidden.retain_grad()

    logits = model.lm_head.output.save()
    logits.sum().backward()

# Access gradient after trace
print(hidden.grad)

Utility Features

Early stopping (skip remaining layers):

with model.trace(prompt) as tracer:
    early_hidden = model.transformer.h[3].output[0].save()
    tracer.stop()  # Skip layers 4+

Shape inference without execution:

with model.scan(prompt):
    shape = model.transformer.h[0].output[0].shape.save()

Remote Execution (NDIF)

For large models, use NDIF's remote infrastructure:

from nnsight import CONFIG
CONFIG.set_default_api_key("YOUR_NDIF_API_KEY")

model = LanguageModel("meta-llama/Llama-3.1-70B")

with model.trace("Hello", remote=True):
    hidden = model.model.layers[-1].output[0].save()

When writing more than one trace against a remote model, see the remote skill for the request-batching and download-pruning patterns that keep latency manageable.

Common Pitfalls

1. Forgetting .save(): Values not saved are lost after tracing
2. Out-of-order access: Access modules in execution order only
3. Using values inside trace: Use .save() and access .value outside the trace