Setting Values

We often not only want to see whats happening during computation, but intervene and edit the flow of information.

In this example, we create a tensor of noise to add to the hidden states. We then add it, use the assigment = operator to update the tensors of .output[0][:] with these new noised values.

from nnsight import LanguageModel
import torch

model = LanguageModel('openai-community/gpt2', device_map='cuda')

with model.trace('The Eiffel Tower is in the city of') as tracer:

    hidden_states_pre = model.transformer.h[-1].output[0].clone().save()

    noise = (0.001**0.5)*torch.randn(hidden_states_pre.shape)

    # model.transformer.h[-1].output = (hidden_states_pre + noise, model.transformer.h[-1].output[1])
    model.transformer.h[-1].output[0][:] = hidden_states_pre + noise

    hidden_states_post = model.transformer.h[-1].output[0].save()

/share/u/caden/.conda/envs/interp/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
You're using a GPT2TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.

We can see the change in the results:

print(hidden_states_pre)
print(hidden_states_post)

tensor([[[ 0.0505, -0.1728, -0.1690,  ..., -1.0096,  0.1280, -1.0687],
         [ 8.7494,  2.9057,  5.3024,  ..., -8.0418,  1.2964, -2.8677],
         [ 0.2960,  4.6686, -3.6642,  ...,  0.2391, -2.6064,  3.2263],
         ...,
         [ 2.1537,  6.8917,  3.8651,  ...,  0.0588, -1.9866,  5.9188],
         [-0.4460,  7.4285, -9.3065,  ...,  2.0528, -2.7946,  0.5556],
         [ 6.6286,  1.7258,  4.7969,  ...,  7.6714,  3.0683,  2.0481]]],
       device='cuda:0', grad_fn=<CloneBackward0>)
tensor([[[-0.0244, -0.1430, -0.1384,  ..., -1.0185,  0.1173, -1.0855],
         [ 8.8112,  2.8995,  5.3000,  ..., -8.0057,  1.3378, -2.8778],
         [ 0.2780,  4.6427, -3.6233,  ...,  0.1378, -2.5877,  3.2030],
         ...,
         [ 2.1590,  6.9278,  3.8584,  ...,  0.0728, -2.0145,  5.9685],
         [-0.4962,  7.4435, -9.3371,  ...,  2.1299, -2.7784,  0.5379],
         [ 6.6941,  1.7279,  4.8105,  ...,  7.6403,  3.0599,  2.0637]]],
       device='cuda:0', grad_fn=<CopySlices>)