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

Early Stopping¶

When you only need activations from the first few layers, there's no reason to run the full forward pass. Call tracer.stop() to abort the run at the current point, saving time and compute.

stop() raises an EarlyStopException that the interleaver treats as a clean early exit (not an error) and swallows. Any module that hasn't executed yet never runs.

Setup¶

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)

Stopping After a Layer¶

tracer.stop() aborts execution at the point the run has currently reached. Everything after that point is skipped. Save what you need before you call it.

In [2]:

with model.trace("The Eiffel Tower is in the city of") as tracer:
    hs = model.transformer.h[5].output[0].save()
    tracer.stop()

print(f"Layer 5 output shape: {hs.shape}")

with model.trace("The Eiffel Tower is in the city of") as tracer: hs = model.transformer.h[5].output[0].save() tracer.stop() print(f"Layer 5 output shape: {hs.shape}")

Layer 5 output shape: torch.Size([10, 768])

The model never ran layers 6–11 or lm_head. Only layers 0–5 executed, which is all we needed to populate hs.

Code After stop() Is Skipped¶

tracer.stop() raises inside the block, so any line after it in the same block never runs — the interventions it would define are never registered.

In [3]:

with model.trace("The Eiffel Tower is in the city of") as tracer:
    hs_early = model.transformer.h[0].output[0].save()
    tracer.stop()

    # This line never runs
    hs_late = model.transformer.h[-1].output[0].save()

print(f"Layer 0: {hs_early.shape}")

try:
    print(hs_late)
except NameError:
    print("hs_late was never defined — stop() halted the block before it")

with model.trace("The Eiffel Tower is in the city of") as tracer: hs_early = model.transformer.h[0].output[0].save() tracer.stop() # This line never runs hs_late = model.transformer.h[-1].output[0].save() print(f"Layer 0: {hs_early.shape}") try: print(hs_late) except NameError: print("hs_late was never defined — stop() halted the block before it")

Layer 0: torch.Size([10, 768])
hs_late was never defined — stop() halted the block before it

Everything after stop() is dead code

tracer.stop() raises immediately, so any .save() calls, module accesses, or other interventions written below it are never registered. Place tracer.stop() after your last intervention.

Two consequences worth knowing before you rely on a stop:

• The run's result is gone. A stopped run never produces one, so a tracer.result.save() after the stop() is unreachable, and moving it into a separate empty tracer.invoke() raises OutOfOrderError: 'result.i0' instead. Save activations, not results.
• Saving tracer.result before the stop() defeats the stop. tracer.result is served when the call returns, so the worker parks on it until the whole forward has run, and the stop() below it only fires afterwards. The tell is the saved object: it carries logits for every position, which a stopped run could not have produced.

Early Stop During Generation¶

During generation, stop() ends the whole run — not just the current step. Collect what you need per step, then bail once a condition is met.

In [4]:

import nnsight

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=20) as tracer:
    tokens = nnsight.save([])
    for step in tracer.iter[:20]:
        tokens.append(model.generator.streamer.output)
        if len(tokens) == 3:
            tracer.stop()

print(f"Steps collected before stopping: {len(tokens)}")
print(f"Step shapes: {[tuple(t.shape) for t in tokens]}")

import nnsight with model.generate("The Eiffel Tower is in the city of", max_new_tokens=20) as tracer: tokens = nnsight.save([]) for step in tracer.iter[:20]: tokens.append(model.generator.streamer.output) if len(tokens) == 3: tracer.stop() print(f"Steps collected before stopping: {len(tokens)}") print(f"Step shapes: {[tuple(t.shape) for t in tokens]}")

Steps collected before stopping: 3
Step shapes: [(1, 10), (1,), (1,)]

Generation would have run 20 steps; stop() cut it off after the third. The first entry is the encoded prompt, then one new token per step.

Performance benefit¶

Early stopping is faster because the model skips every layer after the stop point. How much faster depends on how early you stop and on what else the machine is doing, so treat the number below as one measurement rather than a constant.

In [5]:

import time

# Full forward pass
start = time.perf_counter()
for _ in range(50):
    with model.trace("The Eiffel Tower is in the city of"):
        hs = model.transformer.h[5].output[0].save()
full_time = time.perf_counter() - start

# Early stop after layer 5
start = time.perf_counter()
for _ in range(50):
    with model.trace("The Eiffel Tower is in the city of") as tracer:
        hs = model.transformer.h[5].output[0].save()
        tracer.stop()
stop_time = time.perf_counter() - start

print(f"Full forward pass (50 runs): {full_time:.3f}s")
print(f"Early stop at layer 5 (50 runs): {stop_time:.3f}s")
print(f"Speedup: {full_time / stop_time:.1f}x")

import time # Full forward pass start = time.perf_counter() for _ in range(50): with model.trace("The Eiffel Tower is in the city of"): hs = model.transformer.h[5].output[0].save() full_time = time.perf_counter() - start # Early stop after layer 5 start = time.perf_counter() for _ in range(50): with model.trace("The Eiffel Tower is in the city of") as tracer: hs = model.transformer.h[5].output[0].save() tracer.stop() stop_time = time.perf_counter() - start print(f"Full forward pass (50 runs): {full_time:.3f}s") print(f"Early stop at layer 5 (50 runs): {stop_time:.3f}s") print(f"Speedup: {full_time / stop_time:.1f}x")

Full forward pass (50 runs): 0.375s
Early stop at layer 5 (50 runs): 0.198s
Speedup: 1.9x

When to use early stopping

• Collecting activations from early/middle layers — no need to run the full model
• Training probes or SAEs on intermediate representations — skip the layers you don't need
• Debugging — quickly check a specific layer's output without waiting for the full pass

tracer.stop() is a successful early exit, not an error path. To bypass a single module without aborting the whole forward, use module.skip(...) instead.