Early Stopping

Summary

You can implement more efficient experiments with nnsight by halting a model’s forward pass run early after accessing your values of interest. Calling tracer.stop() will interrup model execution and terminate tracing without going through the rest of the model’s execution.

with model.trace(prompt) as tracer:
   lk_out = model.layer[k].output.save()

   # stop after layer k
   tracer.stop()

When to Use

Early stopping makes targeted experiments more efficient! You can break tracing early if you don’t need to complete the full forward pass or want to control the stopping point of multi-token generation.

It can be particularly helpful for intensive interventions that only need to occur up to a certain layer, like training SAEs for a given layer of a model.

How to Use

Calling tracer.stop() will interrupt tracing immediately at the most recent point of the model’s execution, which will be defined by the latest module being accessed above the stop call.

Let’s try it out! We’ll save the output of layer 5 and then halt the execution.

with model.trace(prompt) as tracer:

  hs = model.transformer.h[5].output.save()

  # stop after layer 5
  tracer.stop()

print(hs)

(tensor([[[ 0.1408, -1.4999,  0.9446,  ..., -1.5783, -0.6470, -0.1584],
         [ 0.8697,  0.3205,  2.2336,  ..., -2.1563, -1.1537, -2.0113],
         [ 1.4215,  2.9198, -0.3065,  ..., -1.1516, -0.2206,  3.1634],
         ...,
         [ 0.0068,  2.2785, -0.7487,  ...,  0.2559,  0.7301,  0.6825],
         [-2.4241,  1.7621, -5.4071,  ...,  3.3570, -0.6088, -0.2757],
         [ 3.7067,  2.9246, -0.8595,  ...,  2.0726, -0.2367,  4.0470]]],
       device='cuda:0', grad_fn=<AddBackward0>),)

Keep in mind that Early Stopping means all subsequently defined interventions will never run.

with model.trace(prompt) as tracer:
  model.transformer.h[0].output

  tracer.stop()

  hs_1 = model.transformer.h[1].output.save()

print("L2 - Output: ", hs_1)

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
/tmp/ipython-input-3130486789.py in <cell line: 0>()
      6   hs_1 = model.transformer.h[1].output.save()
      7
----> 8 print("L2 - Output: ", hs_1)

NameError: name 'hs_1' is not defined

In the example above, the model’s execution was stopped at layer 1, so the output of layer 2 was never defined.

Time Saving

Early Stopping interrupts the model execution early, and hence, renders experiments more time-efficient.

Let’s verify that the model is actually being interrupted and that the time gain is real.

import time

# no stop tun
start_time = time.perf_counter()
with model.trace(prompt):
   hs = model.transformer.h[5].output.save()

no_stop_exc_time = time.perf_counter() - start_time


# early stop run
start_time = time.perf_counter()
with model.trace(prompt) as tracer:
   hs = model.transformer.h[5].output.save()

   # stop after layer 5
   tracer.stop()

stop_exc_time = time.perf_counter() - start_time

print(f"Execution Time w/o Stop: {no_stop_exc_time:.4f}")
print(f"Execution Time w/ Early Stop: {stop_exc_time:.4f}")

Execution Time w/o Stop: 0.0295
Execution Time w/ Early Stop: 0.0150