johnma2006 / mamba-minimal

Simple, minimal implementation of the Mamba SSM in one file of PyTorch.
Apache License 2.0
2.54k stars 188 forks source link

Setting MPS backend causes mamba-minimal to generate gibberish and crash on M1 Mac #15

Open oselivanov opened 8 months ago

oselivanov commented 8 months ago

Trying to run it on MPS backend on Mac M1 Max 64Gb modifying demo.ipynb like this:


Changes to notebook:

device = 'mps'

model = Mamba.from_pretrained(pretrained_model_name).to(device)

input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(device)


Getting gibberish and dead python kernel in dozen of tokens:

for t in generate(model, tokenizer, 'Mamba is the'): print(t, end='')

ocardial goal fibrobl ( triglycer ા judgement extraordinacebook CURIAM

and then it crashes.