Setting

Summary

Setting is how you can intervene on a model and edit information as it flows throughout its computations.

with net.trace(input):

  # setting all output values to 0
  net.layer1.output[:] = 0

  l2_inp = net.layer2.inputs.save()

assert torch.equal(l2_inp == 0)

When to Use

Setting is used when making interventions on the model, like during activation patching.

How to Use

We often not only want to see whats happening during computation, but also to intervene and edit the flow of information in a model. We do this with the setting (=) operation.

There are two ways to perform setting in NNsight: in-place, and not in-place:

• In-place setting operations change the values of the tensor that the targeted output for the module is referencing (which has downstream implications for gradient operations!). In NNsight, in-place setting of a module’s value would look like model.<module_path>.output[:] = new_tensor[:].

• Not in-place setting operations assign a new tensor to the module instead of changing the values of the current tensor. In NNsight, not in-place setting of a module’s value would look like model.<module_path>.output = new_tensor.

This is no different than standard Python syntax, just worth mentioning both are effective. Some objects are immutable (tuples) and would have to be replaced outright instead of in-place.

In-place setting

with tiny_model.trace(input):

    # Save the output before the edit to compare.
    # Notice we apply .clone() before saving as the setting operation is in-place.
    l1_output_before = tiny_model.layer1.output.clone().save()

    # Access the 0th index of the hidden state dimension and set it to 0.
    tiny_model.layer1.output[:,0] = 0

    # Save the output after to see our edit.
    l1_output_after = tiny_model.layer1.output.save()

print("Before:", l1_output_before)
print("After:", l1_output_after)

Before: tensor([[ 0.5857,  0.2823, -0.0138, -0.4004,  0.0133, -0.5596,  0.1553,  0.1798,
         -0.4834,  0.2742]])
After: tensor([[ 0.0000,  0.2823, -0.0138, -0.4004,  0.0133, -0.5596,  0.1553,  0.1798,
         -0.4834,  0.2742]])

Not in-place setting

with tiny_model.trace(input):

    # Save the output before the edit to compare.
    # No need to apply .clone() before saving, as we are not setting in-place.
    l1_output_before = tiny_model.layer1.output.save()

    edited_tensor = tiny_model.layer1.output.clone().save()
    edited_tensor[:, 0] = 0

    # Set the layer1 output equal to the edited tensor
    tiny_model.layer1.output = edited_tensor

    # Save the output after to see our edit.
    l1_output_after = tiny_model.layer1.output.save()

print("Before:", l1_output_before)
print("After:", l1_output_after)

Before: tensor([[ 0.5857,  0.2823, -0.0138, -0.4004,  0.0133, -0.5596,  0.1553,  0.1798,
         -0.4834,  0.2742]])
After: tensor([[ 0.0000,  0.2823, -0.0138, -0.4004,  0.0133, -0.5596,  0.1553,  0.1798,
         -0.4834,  0.2742]])

Order Matters

In NNsight, your intervention code is ran serially alongside the model’s normal execution. Therefore the order you interact with modules must match the order those modules are called in the original underlying model.

For our tiny model above, if I wanted to interact with both layer1 and layer2, I must do layer1 first. In this example I make the mistake of doing layer 2 first, resulting in an OutOfOrderError:

with tiny_model.trace(input):

    l2_output = tiny_model.layer2.output.save()

    l1_output = tiny_model.layer1.output.save()

print("Layer1:", l1_output)
print("Layer2:", l2_output)

---------------------------------------------------------------------------
NNsightException                          Traceback (most recent call last)
/tmp/ipython-input-4160868438.py in <cell line: 0>()
----> 1 with tiny_model.trace(input):
      2
      3     l2_output = tiny_model.layer2.output.save()
      4
      5     l1_output = tiny_model.layer1.output.save()

/usr/local/lib/python3.11/dist-packages/nnsight/intervention/tracing/base.py in __exit__(self, exc_type, exc_val, exc_tb)
    414
    415             # Execute the traced code using the configured backend
--> 416             self.backend(self)
    417
    418             return True

/usr/local/lib/python3.11/dist-packages/nnsight/intervention/backends/execution.py in __call__(self, tracer)
     22         except Exception as e:
     23
---> 24             raise wrap_exception(e, tracer.info) from None
     25         finally:
     26             Globals.exit()

NNsightException:

Traceback (most recent call last):
  File "/tmp/ipython-input-4160868438.py", line 5, in <cell line: 0>
    l1_output = tiny_model.layer1.output.save()
  File "/usr/local/lib/python3.11/dist-packages/nnsight/intervention/envoy.py", line 152, in output
    return self._interleaver.current.request(
  File "/usr/local/lib/python3.11/dist-packages/nnsight/intervention/interleaver.py", line 781, in request
    value = self.send(Events.VALUE, requester)
  File "/usr/local/lib/python3.11/dist-packages/nnsight/intervention/interleaver.py", line 766, in send
    raise response

OutOfOrderError: Value was missed for model.layer1.output.i0. Did you call an Envoy out of order?

Related

• Getting