nnsight icon indicating copy to clipboard operation
nnsight copied to clipboard

Calling `next` after `all` should raise a proper error

Open Butanium opened this issue 10 months ago • 0 comments

Disclaimer: I was not surprised that this fails as it doesn't make sense to use all and next on the same module

from nnsight import LanguageModel
model = LanguageModel("gpt2")
with model.generate(prompt, max_new_tokens=10):
    model.transformer.h.all()
    model.transformer.h.next()

Fails with TypeError: unsupported operand type(s) for +: 'slice' and 'int'

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
File /workspace/clement/repos/representation-structure-comparison/junk.py:4
      2 from nnsight import LanguageModel
      3 model = LanguageModel("gpt2")
----> 4 with model.generate(prompt, max_new_tokens=10):
      5     model.transformer.h.all()
      6     model.transformer.h.next()

File /workspace/clement/repos/representation-structure-comparison/.env/lib/python3.10/site-packages/nnsight/intervention/contexts/interleaving.py:96, in InterleavingTracer.__exit__(self, exc_type, exc_val, exc_tb)
     92     self.invoker.__exit__(None, None, None)
     94 self._model._envoy._reset()
---> 96 super().__exit__(exc_type, exc_val, exc_tb)

File /workspace/clement/repos/representation-structure-comparison/.env/lib/python3.10/site-packages/nnsight/tracing/contexts/tracer.py:25, in Tracer.__exit__(self, exc_type, exc_val, exc_tb)
     21 from .globals import GlobalTracingContext
     23 GlobalTracingContext.try_deregister(self)
---> 25 return super().__exit__(exc_type, exc_val, exc_tb)

File /workspace/clement/repos/representation-structure-comparison/.env/lib/python3.10/site-packages/nnsight/tracing/contexts/base.py:72, in Context.__exit__(self, exc_type, exc_val, exc_tb)
     69 graph = self.graph.stack.pop()
     71 if isinstance(exc_val, BaseException):
---> 72     raise exc_val
     74 self.add(graph.stack[-1], graph, *self.args, **self.kwargs)
     76 if self.backend is not None:

File /workspace/clement/repos/representation-structure-comparison/junk.py:6
      4 with model.generate(prompt, max_new_tokens=10):
      5     model.transformer.h.all()
----> 6     model.transformer.h.next()

File /workspace/clement/repos/representation-structure-comparison/.env/lib/python3.10/site-packages/nnsight/intervention/envoy.py:251, in Envoy.next(self, increment)
    241 def next(self, increment: int = 1) -> Envoy:
    242     """By default, this modules inputs and outputs only refer to the first time its called. Use `.next()`to select which iteration .input an .output refer to.
    243 
    244     Args:
   (...)
    248         Envoy: Self.
    249     """
--> 251     return self.iter[self._iteration_stack[-1] + increment].__enter__()

TypeError: unsupported operand type(s) for +: 'slice' and 'int'

For newcomer, it'd be probably better if this raise a proper error message indicating that next should not be used on a module affected by all

Butanium avatar Mar 11 '25 06:03 Butanium