Add codestral mamba2
Author: molbap
Pull request: yes
Labels: name: run-slow

# What does this PR do?

This PR adds mamba2 to support mistral's mamba-codestral release.

- [x] Make sure weights are correctly loaded
- [x] Call correct cuda kernel in case of self.training, like mamba_inner_fn for mamba1
- [x] Cuda forward actually supports outputting states - we can use it for caching
- [x] Implement slow_forward as well
- [x] fix quality of slow forward generations
- [x] remove einops
- [x] Really remove all einops
- [x] Figure out why cuda kernels are actually slow? edit: it's unclear but an issue raised in the original repo https://github.com/state-spaces/mamba/issues/389#issuecomment-2171755306
- [x] add tests
- [x] Fix batched generation by tuning out states (thanks @vasqu !), edit: pointed out in https://github.com/state-spaces/mamba/issues/66#issuecomment-1863563829, affects mamba-1 as well likely
- [x] add more integration tests (batched generation)

cc @ArthurZucker 

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_32080). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@molbap I believe it will be more efficient to build on top of the current Mamba2 PR https://github.com/huggingface/transformers/pull/32027 to speed up the integration process!

Hey both! We tried your code and could not get it to work, mostly the shape of the output final_state is wrong (ends with 1024 vs 128 that we need). We are probably gonna ship it as is, but a patch will be possible if we can ship the ssimd in it! 

@ArthurZucker I've just seen that it's using groups, so the code doesn't work out of the box as it would need appropriate resizing of B and C for example (ngroups -> nheads by repetition). I did not account for that as the original models never used anything other than group=1. Just tested it locally and it seems to work with those slight adjustments on B and C.  

I can't check in on the final_state shape without seeing the code so maybe I can help later when you guys ship it and there is a base to work on :D  

Hey both again, thanks for the comment @vasqu ! indeed it seems that it needs to be slightly modified - a `B.repeat(1, 1, self.num_heads // self.num_groups, 1)` seems promising. Are you verifying with a .generate() what the model is outputting? We'll ship this soon as Arthur mentioned and credit your code/include you as coauthors of course :hugs: 

@molbap I've been using my [mamba2-torch](https://github.com/vasqu/mamba2-torch) repo to check modifications on the original `ssd_minimal` script to still work correctly (i.e. with variable seq_len, D skip connection, and just now with the n_groups >= 1). You can check it out over [here](https://github.com/vasqu/mamba2-torch/blob/main/tests/TestSSDMinimal.py). It does use einops though which I've changed in our PR to be pure torch.

I changed up generations - the cuda kernel calls + the naive call have slightly different generations but both are very reasonable and seem similar to the original implementation. Missing some tests + remove the triple einsum contraction to create Ydiag and should be good to go!

Alright, I've added a non-einsum equivalent of the tensor contraction paths for both the diag and off-diag terms that seem to give good results, supports batched inputs as well. 

@molbap Did generation improve with the new cache path? Just for my interest 😄 

@vasqu Haha, unfortunately it didn't it seems - there's still a difference in quality between the two paths. But I think it's reasonable for now, would need to do proper benchmarks. 

Just saw the "Figure out why cuda kernels are actually slow?"

I've also wondered about that when i worked on mamba2 and here's an issue of mine I used to keep track: https://github.com/vasqu/mamba2-torch/issues/1#issue-2349175830

It seems like the triton kernels cause a significant cpu overhead on the first iteration(s): https://github.com/state-spaces/mamba/issues/389#issuecomment-2171755306

Other suggestions: big model and/or compiling https://github.com/state-spaces/mamba/issues/355#issuecomment-2147597457 (but that breaks in a lot of cases which happened in my case too)

Yes! Been reading that one too and couldn't get the compile to improve performances (ie crashes). The warmup run definitely helps as well. So this full torch path is currently faster. 

I added some tests and did last touches on modeling - Batched generation is a bit weird, I think there is leaking of states within a batch and not sure this is fixed in the original repo either so keeping it as is.  

Late night me, regarding batched generation: Do the padding sides differ much (left vs right)? I'd assume that we have leaky information flowing through either way. This affects both mamba 1+2. 

When using left padding we would pad with the eos token and subsequently with a non-zero input which would influence causal directions (left to right, repeated eos information is given to the SSM and conv1d). In attention a simple attention mask would take care of it for the core softmax operations, for mamba there does not exist such thing yet. I'd suggest using the original attention mask to zero the padding tokens before entering each mamba block (zero input and zero initial hidden state would result in repeated zeros until encountering the first non zero input). It has been similarly suggested back then for the original mamba model: https://github.com/state-spaces/mamba/issues/66#issuecomment-1863563829

Right padding would probably need keeping track of the sequence indices which recent versions of mamba and conv1d support. Without recent versions it would probably get less "elegant". 

Indeed, that's likely the issue, thanks @vasqu ! I tried tuning out the hidden states but didn't think of doing it in two places :face_with_peeking_eye: It seems that fixes the batched generation, I'll push that soon

@molbap Does that also work for right padding? Reassigning zeros would probably spit out garbage imo (without retrieving the correct last values). Not too familiar with right padding tbh.

Maybe raise a warning and do not interact with the attention mask if right padding, e.g. by setting the mask to None? 

I don't think it does right now for right padding no, but it should work in all other cases now at least!

The `.md` file should include your findings about speed performances etc

Addressed comments and fixed some more tests, also tried a simple PEFT script (almost same as mamba 1) which seems to run :grinning: added it to the .md. 

cc @ArthurZucker , the only test failing is the slow test because runner doesn't have access

Looking really great! :fire: Thanks for the work

let me have a final look and we should be good to merge

Before CI is fixed, the 3 integration tests pass locally

With torch forward:
![image](https://github.com/user-attachments/assets/22030f78-62a9-4c31-8468-f56fb3ec71f1)
With CUDA kernels:
![image](https://github.com/user-attachments/assets/90b5fcb0-49f8-48d4-b56e-71689710956f)



Super excited for this merge!
@vasqu Should we close https://github.com/huggingface/transformers/pull/32027 now?


@Adibvafa Sry, it kinda drowned in my notifications! I'll close it for now, but if needed, it can also be reopened at any time.