Unpacking tuples of tensors failing

General
1

Hi there,

I’m running into an error with torch.topk():

import nnsight
import torch

gpt2 = nnsight.LanguageModel(
    "openai-community/gpt2", device_map="cuda:0", dispatch=True
)

gpt2.requires_grad_(False)

with gpt2.trace("test"):
    actvs = gpt2.transformer.h[0].mlp.output
    vals, idxs = torch.topk(actvs, k=3)

With 0.4.8, I’m getting the following error:

Traceback (most recent call last):
  File "/var/metrics/g/crosslayer-transcoder/nnsight_error.py", line 10, in <module>
    with gpt2.trace("test"):
         ^^^^^^^^^^^^^^^^^^
  File "/var/metrics/g/crosslayer-transcoder/.venv/lib/python3.12/site-packages/nnsight/intervention/contexts/interleaving.py", line 96, in __exit__
    super().__exit__(exc_type, exc_val, exc_tb)
  File "/var/metrics/g/crosslayer-transcoder/.venv/lib/python3.12/site-packages/nnsight/tracing/contexts/tracer.py", line 25, in __exit__
    return super().__exit__(exc_type, exc_val, exc_tb)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/var/metrics/g/crosslayer-transcoder/.venv/lib/python3.12/site-packages/nnsight/tracing/contexts/base.py", line 72, in __exit__
    raise exc_val
  File "/var/metrics/g/crosslayer-transcoder/nnsight_error.py", line 12, in <module>
    vals, idxs = torch.topk(actvs, k=3)
    ^^^^^^^^^^
  File "/var/metrics/g/crosslayer-transcoder/.venv/lib/python3.12/site-packages/nnsight/tracing/graph/proxy.py", line 293, in __iter__
    return iterator.handle_proxy(inspect.currentframe().f_back, self)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/var/metrics/g/crosslayer-transcoder/.venv/lib/python3.12/site-packages/nnsight/tracing/hacks/iterator.py", line 103, in handle_proxy
    end = frame.f_lineno + (for_node.end_lineno - for_node.lineno)
                            ^^^^^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'end_lineno'

I couldn’t test it yet with 0.5.0.devX because I get segmentation faults with all devX

2

Update

I found similar errors, for example with

a, b, c = tensor.shape
# or
5 == a

However, thanks to JadenFK, I got 0.5.0.dev7running now (after the segmentation fault issue with the previous versions).

And with that version, it works! :partying_face:

1 Like