Info
Last Execution: 2026-09-03
| Package | Version |
|---|---|
| nnsight | 0.8.0 |
| Python | 3.12.13 |
| torch | 2.13.0+cu126 |
| transformers | 5.15.0.dev0 |
Skipping Modules¶
Use .skip() to bypass a module's forward pass entirely, substituting a value you provide as its output. When the model is about to run the module, it won't — your replacement is used instead, and none of the module's inner computation happens. This lets you ablate a component, route around a layer, or inject an externally-computed activation (e.g. a reconstruction from an SAE) at a specific point.
from nnsight.modeling.transformers import TransformersModel
model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)
Skipping a Module¶
Call .skip(replacement) on any module to bypass its forward pass. The replacement is used as the module's output instead. Here we skip transformer block 8 and feed the previous block's output through in its place — block 8's forward never runs, so its contribution to the residual stream is dropped and the top prediction changes.
# Normal prediction
with model.trace("The Eiffel Tower is in the city of"):
normal_logits = model.lm_head.output.save()
# Skip transformer block 8 — use block 7's output as its output instead
with model.trace("The Eiffel Tower is in the city of"):
model.transformer.h[8].skip(model.transformer.h[7].output)
skipped_logits = model.lm_head.output.save()
print(f"Normal: {model.tokenizer.decode(normal_logits[0, -1].argmax(dim=-1))}")
print(f"Skipped: {model.tokenizer.decode(skipped_logits[0, -1].argmax(dim=-1))}")
Normal: Paris Skipped: London
The magnitude of the effect depends on how much you drop: skipping a whole block flips the top token here, whereas bypassing a single small sub-module (e.g. one block's MLP) often leaves the top-1 prediction unchanged because its contribution is small.
Pass-through with a Module's Own Input¶
Skipping a module with its own .input turns it into a pass-through: the input is handed straight to the output and the module's forward is bypassed. For a transformer block this is a clean way to drop the block's contribution entirely.
with model.trace("The Eiffel Tower is in the city of"):
# Layer 6 becomes a no-op: its input passes straight through
model.transformer.h[6].skip(model.transformer.h[6].input)
passthrough_logits = model.lm_head.output.save()
print(f"Layer 6 passed through: {model.tokenizer.decode(passthrough_logits[0, -1].argmax(dim=-1))}")
Layer 6 passed through: London
Feeding a module its own input is not always valid
A pass-through only works when a module's input has the same format as its output — as it does for a residual-stream transformer block, whose input and output are both the (batch, seq, hidden) hidden state. It is not universally valid: many modules return something shaped differently from what they receive (a projection changes the hidden size; an attention sub-module may return a tuple; a module may take several arguments but return one tensor). The replacement must match what the module would normally return, not what it takes as input — so when input and output formats differ, module.skip(module.input) will error downstream. Check the module's real .output type and shape first if you're unsure.
Skipping Multiple Layers¶
You can skip a range of layers in a loop. Pass the output of the last non-skipped layer as the replacement for each skipped layer.
with model.trace("The Eiffel Tower is in the city of"):
# Skip layers 3 through 8, reusing layer 2's output for each
replacement = model.transformer.h[2].output
for i in range(3, 9):
model.transformer.h[i].skip(replacement)
logits = model.lm_head.output.save()
print(f"Skipped layers 3-8: {model.tokenizer.decode(logits[0, -1].argmax(dim=-1))}")
Skipped layers 3-8: the
A replacement has to match the output in structure, shape, dtype and device
The replacement stands in for what the module would have returned, so it has to match that value on all four counts. What a module returns is model- and transformers-version-dependent: some hand back a plain tensor (batch, seq, hidden), others a tuple or another structure. Feeding one module's .output to another's .skip() drops in cleanly precisely because both share that return type.
A mismatch is caught by the model, not by nnsight, so it arrives as a bare torch error from inside the next forward with no mention of skip or of the module you skipped:
| Replacement | What surfaces |
|---|---|
x.double() |
RuntimeError: expected scalar type Double but found Float |
x.half() |
RuntimeError: expected scalar type Half but found Float |
x.cpu() |
RuntimeError: Expected all tensors to be on the same device, but got weight is on cuda:0, different from other tensors on cpu |
(x,) around a tensor output |
TypeError: layer_norm(): argument 'input' (position 1) must be Tensor, not tuple |
x[:, :3, :] |
RuntimeError: shape '[-1, 10, 768]' is invalid for input of size 2304 |
When a traced model errors inside a forward you did not intervene on, the skip above it is the first place to look. Read the module's .output first and check its type, shape, dtype and device before choosing a replacement.
Measuring Layer Importance¶
Skip each layer one at a time and check how the prediction changes — a simple way to measure which layers matter most for a given prompt.
prompt = "The Eiffel Tower is in the city of"
with model.trace(prompt):
baseline = model.lm_head.output[0, -1].argmax(dim=-1).save()
baseline_token = model.tokenizer.decode(baseline)
print(f"Baseline: {baseline_token}\n")
for layer_idx in range(1, 12):
with model.trace(prompt):
model.transformer.h[layer_idx].skip(model.transformer.h[layer_idx - 1].output)
pred = model.lm_head.output[0, -1].argmax(dim=-1).save()
token = model.tokenizer.decode(pred)
changed = " <- changed!" if token != baseline_token else ""
print(f"Skip layer {layer_idx:2d}: {token}{changed}")
Baseline: Paris Skip layer 1: London <- changed! Skip layer 2: London <- changed! Skip layer 3: the <- changed! Skip layer 4: Paris Skip layer 5: Paris Skip layer 6: London <- changed! Skip layer 7: Paris Skip layer 8: London <- changed! Skip layer 9: London <- changed! Skip layer 10: London <- changed! Skip layer 11: Paris
Gotchas
- A skipped module's inner ops are unreachable. Since the forward never runs, requesting a skipped module's sub-modules or
.sourceoperations raises an out-of-order error — they never execute. skiponly works inside an active trace, and a skip is one-shot per module call. During multi-token generation, each step needs its own skip (seetracer.iter[...]).- Across batched invokes, a
.skip()must cover every row. A shared forward can't run for only some rows — skip the module in every invoke or none, and each invoke's replacement fills its own rows. - Use
tracer.stop()to abort the whole forward —skiponly bypasses one module.
When to use skip
- Ablation studies — measure the causal effect of removing a layer or sub-module
- Layer importance — identify which layers are critical for specific predictions
- Model splicing — replace a module's computation with an alternative output (e.g. an SAE reconstruction)
- Performance — a skipped module's forward never runs, so substituting a cheap replacement avoids computing it at all