Repository navigation
Feature Request: Support Codestral Mamba #8519
Description
Feature Description
New 7B coding model just released by Mistral.
- Blog Post: https://mistral.ai/news/codestral-mamba/
- HF: https://huggingface.co/mistralai/mamba-codestral-7B-v0.1
Motivation
Seems to perform very well, especially for a 7B model:
Possible Implementation
An extension to #7727?
Reacted by sorasoras, ice0, Kozuch, Duc-Viet Hoang, Nathan Simons, Julian, JohnClaw, Sam, Mirek, randoentity and 74 more
Reacted by Thomas Lorreyte, Sam, lin72h, suncloudsmoon, John Doe, Benjamin Leblanc, Johan, crocodile13, fixerivan, G30 and 3 more
Activity
HanClinto commentedon Jul 16, 2024
I love the shout-out in the linked blog post!
That's a really nice nod -- love to see it!
theo77186 commentedon Jul 16, 2024
#7727 should cover for this model, but with untied embeddings unlike the other Mamba2 models.
timlacroix commentedon Jul 17, 2024
FYI, there is an "ngroups" param that changes how layer norm is done : https://github.com/state-spaces/mamba/blob/c0a00bd1808881831ddf43206c69362d4df90cf7/mamba_ssm/modules/mamba2.py#L47
We use ngroups=8. If you forget it or try with ngroups = 1 you'll have a bad time.
Good luck !
ggerganov commentedon Jul 17, 2024
After we merge #8526 we should try to add full support for this model. cc @compilade
0wwafa commentedon Jul 17, 2024
I'd love this.
txhno commentedon Jul 18, 2024
thanks!
fredconex commentedon Jul 29, 2024
Hey guys, any progress on ETA for it?
rmusser01 commentedon Aug 17, 2024
For anyone else, seems this is waiting on
#8526
which is waiting on
#8980 -> which is waiting on review(?).
compilade commentedon Aug 17, 2024
Some progress report: I have a local branch (not yet public) on top of #8526 in which I've started implementing the graph for Mamba-2. The conv step is very similar to Mamba-1, and I've started to implement the SSM step and will continue in the next days. It's not in a usable state yet.
I'm starting by implementing the fully recurrent mode of Mamba-2 (which is very similar to Mamba-1) (and which is described in Section 3.4.1).
But I'm still evaluating how the block decomposition would fit within how
src/llama.cppmanages batches and/or if the chunk size should be dynamic. It seems like to fully benefit from Section 6, the chunks should be smaller than the batch size, but not too small, at which point directly doing the recurrence is the same. Since theggmlcompute graph nodes should keep the same structure between batches and that the block decomposition will likely have too much overhead for small batches, it's easier to simply go with the linear recurrence with something likeggml_ssm_scanat first.For the ETA, I'll try to get it working before the end of August, but no promises.
(and BTW @rmusser01, #8980 is waiting on #8526, not the other way around, at least I think?)
compilade commentedon Aug 19, 2024
Okay, the fully recurrent mode works for
Mamba-2! (for the curious, see this branch: https://github.com/compilade/llama.cpp/tree/compilade/mamba2) I'll open a PR soon (in the next days; still need to clean up some things).Note that Mamba-Codestral-7B-v0.1 cannot be converted as-is; either use https://huggingface.co/mistralai/Mamba-Codestral-7B-v0.1/discussions/9, or rename
consolidated.safetensorstomodel.safetensors,tokenizer.model.v3totokenizer.model, andparams.jsontoconfig.json. Then, inconfig.json, the line"architectures": ["Mamba2ForCausalLM"],needs to be added (if missing).The state in Mamba-2 is bigger than I thought; Mamba-Codestral-7B-v0.1 takes
263.5 MiB(inF32) per sequence (e.g. with-np 1), compared to38 MiBfor Falcon-Mamba-7B (which is based on Mamba-1). But that remains constant whatever the context size.A big downside right now with recurrent models in
llama.cppis the lack of state rollback (which is implemented through state checkpoints in #7531, but needs to be re-adapted to #8526), so the prompt will be reprocessed a lot if usingllama-server. I think usingllama-cliin conversation mode does not have this problem, however (or maybe only the bare interactive mode with--in-prefixand--in-suffix, not sure).The implementation is CPU-only, but uses SIMD for the SSM scan, so even though the state is bigger than for Mamba-1 models, in my tests, the speed of
Mamba-2-130Mis similar or better thanMamba-130M(but still not that fast compared to transformer-based models with an empty context).The speed of Mamba-2 models seems comparable to Transformer-based models when the latter have 2k to 4k tokens in their context.
Just making sure expectations are not too far from reality.
github-actions commentedon Oct 4, 2024
This issue was closed because it has been inactive for 14 days since being marked as stale.
11 remaining items
massi-ang commentedon Feb 14, 2025
We need this. Any news on when this will be available?
moonray commentedon Mar 28, 2025
Would love to get an update on this.
github-actions commentedon Jun 16, 2025
This issue was closed because it has been inactive for 14 days since being marked as stale.
compilade commentedon Aug 28, 2025
Heads up that #15625 fixes a problem in the implementation of
SSM_SCAN, which makes this model (Mamba-Codestral-7B-v0.1) better than it was when initially implemented here.So if you had some initial impressions from this model from when it was first implemented here, it might be relevant to revisit this model after the recent changes.