Yeah, I’m trying this and getting this error
from nnsight import LanguageModel
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token
model_nn = LanguageModel(model, tokenizer=tokenizer)
prefix = "I went to the kitchen"
prefix_inputs = tokenizer(prefix, return_tensors="pt")
new_text = " and I got a sandwich"
with torch.no_grad():
prefix_outputs = model(**prefix_inputs, use_cache=True)
with torch.no_grad():
with model_nn.trace(new_text, past_key_values=prefix_outputs.past_key_values, use_cache=True) as trace:
h = model_nn.transformer.h[-1].output[0].save()
Traceback (most recent call last):
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/nnsight/tracing/graph/node.py", line 289, in execute
self.target.execute(self)
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/nnsight/intervention/contexts/interleaving.py", line 161, in execute
graph.model.interleave(interleaver, *invoker_args, fn=method,**kwargs, **invoker_kwargs)
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/nnsight/modeling/mixins/meta.py", line 52, in interleave
return super().interleave(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/nnsight/intervention/base.py", line 341, in interleave
with interleaver:
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/nnsight/intervention/interleaver.py", line 129, in __exit__
raise exc_val
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/nnsight/intervention/base.py", line 342, in interleave
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/nnsight/modeling/language.py", line 297, in _execute
return self._model(
^^^^^^^^^^^^
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1603, in _call_impl
result = forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 1062, in forward
transformer_outputs = self.transformer(
^^^^^^^^^^^^^^^^^
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1603, in _call_impl
result = forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 829, in forward
attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/transformers/modeling_attn_mask_utils.py", line 395, in _prepare_4d_causal_attention_mask_for_sdpa
expanded_4d_mask = attn_mask_converter.to_4d(
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/gustaw/Documents/concept_vectors/.venv/lib/python3.12/site-packages/transformers/modeling_attn_mask_utils.py", line 139, in to_4d
expanded_attn_mask = causal_4d_mask.masked_fill(expanded_attn_mask.bool(), torch.finfo(dtype).min)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: The size of tensor a (5) must match the size of tensor b (10) at non-singleton dimension 3
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
NNsightError: The size of tensor a (5) must match the size of tensor b (10) at non-singleton dimension 3```