# Nnsight Basics

> Core nnsight concepts for neural network interpretability. Use when setting up models, tracing activations, saving values, or making basic interventions on model internals.

- Skill: `ndif-team/nnsight-basics` (Agent Skill)
- Install (CLI): `npx skillmds@latest add ndif-team/nnsight-basics`
- Raw SKILL.md: https://api.skillmd.com/api/skills/ndif-team/nnsight-basics/raw
- Safety review: pending (external: skill-scanner PASS, skillspector PASS)
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: AI & ML
- Author: ndif-team (https://skillmd.com/u/ndif-team)
- Updated: 2026-09-10
- Page: https://skillmd.com/skills/ndif-team/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

```python
pip install nnsight
pip install torch transformers
```

## Loading Models

**For language models (recommended):**

```python
from nnsight import LanguageModel

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

**For any PyTorch model:**

```python
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.

```python
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

| Property | Returns |
| -------- | ------- |
| `.output` | Module's return value |
| `.input` | First positional argument |
| `.inputs` | All arguments (tuple + kwargs dict) |

```python
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:

```python
# 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`:

```python
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.

```python
# 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:

```python
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:**

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

**Add noise:**

```python
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:**

```python
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:**

```python
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:

```python
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.

```python
# 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:

```python
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:

```python
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:

```python
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:**

```python
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:**

```python
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):

```python
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):**

```python
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):**

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

**Shape inference without execution:**

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

## Remote Execution (NDIF)

For large models, use NDIF's remote infrastructure:

```python
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