search

state-spaces / mamba

Mamba SSM architecture
Apache License 2.0
13.36k stars 1.13k forks source link

Export mamba models #626

Open bhack opened 1 week ago

bhack commented 1 week ago

Can you add a small example on how to export mamba models?

See https://github.com/pytorch/pytorch/issues/130150

https://pytorch.org/tutorials/intermediate/torch_export_tutorial.html#custom-ops

Thanks

vasqu commented 1 week ago

Iirc, I think the team at meta showed it in a pretty nice way over here: https://github.com/facebookresearch/lingua/tree/main/apps/mamba/component

Edit: at least for Mamba2

Edit 2: I was too quick, i mistook it for compiling 🤦 ignore this

bhack commented 1 week ago

I am asking this cause https://github.com/huggingface/transformers/issues/33839.

But yes after exporting the model it could be also compiled so a complete example of export+compile it will be very helpful.

bhack commented 5 days ago

@mathuvu Do you plan to contribute your https://github.com/facebookresearch/lingua/tree/main/apps/mamba/component upstream?