Possible of implementing mamba ssm

[tikikun]

Prerequisites

Please answer the following questions for yourself before submitting an issue.

• I am running the latest code. Development is very rapid so there are no tagged versions as of now.
• I carefully followed the README.md.
• I searched using keywords relevant to my issue to make sure that I am creating a new issue that is not already open (or closed).
• I reviewed the Discussions, and have a new bug or useful enhancement to share.

Feature Description

As a frontier project in edge inferencing, mamba ssm can have huge impact if the quality of the arch is as good as claimed. It will boost edge inference to another level.

Motivation

https://arxiv.org/abs/2312.00752

[C0deMunk33]

has anyone picked this up? I'm excited and would like to use it in my current llama.cpp stack

[paryska99]

I hope these Mamba models are actually going to scale well with bigger parameter counts. Still very exciting to see such fast paced development in AI space.

[bachittle]

Would this be out of scope for llama.cpp? Perhaps it would be more beneficial to make a new library and add as an example to ggml.

[ekg]

It would be good to be able to host and use mamba models in gguf format with all the quantization approaches that they can be linked with. And people happily build on llama.cpp. I'd be up for working on an implementation, but need some confirmation that adding it here makes sense. But adding a new model type might be a big pain so moving this to another project is possibly cleaner.

[ggerganov]

I don't think anything prevents from adding this model to llama.cpp, though I don't have a good estimate of the amount of effort that it would require to do so. I suppose we will give it a shot sometime in the future after dealing with some of the higher-priority things. Unless of course there is help from the community in the meantime.

[rahuldshetty]

If anybody is picking up, they can have a look at an implementation from @johnma2006: https://github.com/johnma2006/mamba-minimal/tree/master

[LegallyCoder]

https://github.com/LegallyCoder/mamba-hf
My implementation is here.

[AlbertMarashi]

Interested in this. Would like to use in LM Studio

[LegallyCoder]

Interested in this. Would like to use in LM Studio

Maybe, but it's not fully finished, it needs to be developed. But nowly almost all functions support, except some, I'm trying to fix them. Sure if you tell me

[kroggen]

Regarding quantization, this is from the README:

Precision

Our models were trained using PyTorch AMP for mixed precision. AMP keeps model parameters in float32 and casts to half precision when necessary. On the other hand, other frameworks like DeepSpeed store parameters in float16 and upcasts when necessary (e.g. for optimizer accumulation).

We've observed that higher precision for the main model parameters may be necessary, because SSMs are sensitive to their recurrent dynamics. If you are experiencing instabilities, as a first step please try a framework storing parameters in fp32 (such as AMP).

[kroggen]

The 2 implementations above only have the convolution mode (good for training) and lack the recurrent mode (good for inference)

The original implementation has support for both modes. You can pay attention to the inference_params and _decoding_cache used in the recurrent mode, as well as mamba_simple.py

This is my fork that works on CPU (the original requires CUDA and a GPU)

There is an implementation for tinygrad that also has support for the recurrent mode

[cztomsik]

Forgive me if I misunderstood the whole mamba paper but recurrent mode means it's like RNN (like RWKV) and that also means it has constant memory usage and therefore it cannot see "backwards".

Which means:

• you need to rephrase tasks in prefix way (not a big issue)
• it's way more sensitive to sampling (randomness can help but it can also damage the state significantly)
• it might derail chat conversation in a totally useless way, because no matter what you see on the screen, it might not be in the "state" anymore.

I am not trying to bash the model, I think RNNs are great for many things, I'd just like to understand limitations correctly.