The mambalens from phylliida

mambalens's People

Contributors

Stargazers

Watchers

mambalens's Issues

Error on compiling InputDependentHookPoint with fullgraph=True

Compiling a network consisting of InputDependentHookPoint with fullgraph=True results in the following error:

  File "/opt/conda/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2312, in inline_call_
    raise ArgsMismatchError(  # noqa: TRY200
torch._dynamo.exc.ArgsMismatchError: got an unexpected keyword argument 'postfix'.
  func = 'forward' /opt/venv/lib/python3.10/site-packages/transformer_lens/hook_points.py:127, args = [<class 'mamba_lens.input_dependent_hooks.InputDependentHookPoint'>, <class 'torch.Tensor'>], kwargs = {'postfix': ConstantVariable(str)}

from user code:
   File "/opt/venv/lib/python3.10/site-packages/mamba_lens/HookedMamba.py", line 901, in forward
    h        = self.hook_h(h, postfix=postfix) # [B,E,N]
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)

The below minimal code can be used to reproduce the error

import os
import torch
from mamba_lens import HookedMamba

os.environ["TOKENIZERS_PARALLELISM"] = "false"
model = HookedMamba.from_pretrained("state-spaces/mamba-370m", device='cuda')
input_tensor = model.tokenizer("Hello World", return_tensors="pt").to("cuda")
embs = model.embedding(input_tensor["input_ids"])
compiled_first_block = torch.compile(model.blocks[0], fullgraph=True)
block_output = compiled_first_block(embs)
print(block_output.shape)