search

ndif-team / nnsight

The nnsight package enables interpreting and manipulating the internals of deep learned models.
https://nnsight.net/
MIT License
356 stars 34 forks source link

Envoy input access syntactic change #205

Closed AdamBelfki3 closed 3 weeks ago

AdamBelfki3 commented 3 weeks ago

NNsight syntactic change!

Envoy.input now points directly to the first positional argument of a module, i.e. the first tensor input.

from nnsight import LanguageModel

model = LanguageModel("openai-community/gpt2", device_map="auto")

with model.trace("Hello World"):
        hs = model.transformer.h[6].input.save()

print("Tensor: ", hs)
"""
Tensor: tensor([[[ 0.9839, -2.3671,  0.9822,  ..., -1.1355, -0.6406, -1.1865],
         [-0.6150, -0.0651,  2.2578,  ..., -0.2923, -1.8130,  2.8994]]],
       device='mps:0', grad_fn=<AddBackward0>)
"""

If you still wish to access to full module input, you can use Envoy.inputs:

from nnsight import LanguageModel

model = LanguageModel("openai-community/gpt2", device_map="auto")

with model.trace("Hello World"):
        hs = model.transformer.h[6].inputs.save()

print("Input: ", hs)
"""
Input: ((tensor([[[ 0.9839, -2.3671,  0.9822,  ..., -1.1355, -0.6406, -1.1865],
         [-0.6150, -0.0651,  2.2578,  ..., -0.2923, -1.8130,  2.8994]]],
       device='mps:0', grad_fn=<AddBackward0>),), {'layer_past': None, 'attention_mask': None, 'head_mask': None, 'encoder_hidden_states': None, 'encoder_attention_mask': None, 'use_cache': True, 'output_attentions': False})
"""