[`Mamba2`] Fix caching, slow path, and multi-gpu
Author: vasqu
Pull request: yes
Labels: name: run-slow

name: State space models

# What does this PR do?
Kind of a follow-up to #34901 as there are some issues in the current code:
- Caching
- Dt clamping
- Following Mamba1 standards a bit closer
- Multi-Gpu fix when caching
- Slow path fix (supersedes #34901)

Fixes #33567
Fixes #34817

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @ylacombe, @eustlb
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->
@molbap @ArthurZucker 

Comments:
Integration tests will probably need an update but I don't have a GPU for the 7B atm.

Edit: If you could update these integration tests, then gladly :D especially since I'm on vacay very soon

Hey @vasqu thanks! Taking a look in a min

also I added the slow label - feel free to launch a commit with message "[run-slow] mamba2" so we can trigger the slow CI! that way we make sure multi-gpu is indeed fixed

@molbap Yup, added an empty commit - will get to the comments/review a bit later :saluting_face: 

(I could expect some failures on the integration tests, not sure let's see)

Attempt 2 at multi gpu, at least a different error :p 

Things that remain:
- [ ] Mulit-GPU *sigh*
- [ ] Supersede slow path fix with same tests included?
- [ ] Compile compatibility (future)
- [ ] Refactor some stuff (future)

Otherwise, ready to go @molbap 

Edit: Hub seems to have some unrelated issues 

Could you trigger the slow runs @molbap ? I added the slow path fix here and closed the old PR. The only thing that didn't work before was the multi GPU caching so let's see how it goes this time.

I'm on vacay for a while, so feel free to update / correct some stuff if you want to. Will be back after xmas, so I guess early merry xmas :)

Got it, will update if needed! Enjoy the holidays @vasqu !

Thanks for the effort @vasqu! Just wanted to quick check in to understand the status of current HF implementation of Mamba-2.

### `cuda_kernels_forward`
- The current implementation correct, since most of the proposed changes are style changes.

### `torch_forward`
- The current implementation is incorrect, and the PR fixes two things.
1. Reduction dimension.
2. `dt` clamping.

- There are some updates to masking, but not sure when/whether these are used.

Are these statements correct? Thanks in advance!

Hey @HanGuo97, you understood most of the changes imo. I guess only things I'd add on would be the caching issues specifically:
- cuda path wrongly initiates the conv cache
- torch path had some issues with dimensions when caching / inferring (I think you mean this as well)
- overall cleanup of how inference is prepared as there was a mixup in regards to non-cached inference
- multi-gpu device management caused issues (when interacting with the cache)

The masks were there before but a bit of a refactor tbh. They're important to ensure that batched inference works correctly - that's something that is incorrect in the original mamba(2) implemenation (except you use varlen implementations, i.e. those that use cu_seq_lens).

Lmk if something's unclear! 

Thank you for the clarifications! 

Is the current PR mostly complete, or can we expect more significant fixes? I’m asking because I’m refactoring the implementation for other reasons and prefer to base it on a more stable version.

On a separate note, do you have any suggestions for testing the implementation’s correctness? I was considering comparing it with the original Mamba-2 implementation, but I’m unsure if I might overlook something.

No problem! I don't think that there will be any more significant changes. So, this should be a good ground to work on.

The original implementation is well made as is. I only have my gripes when using batches. Depending on your refactor it might easier to look at earlier versions without cu_seq as it adds a bit of complexity. Otherwise, kind of self-promo but I've also written a mamba2 version back before it came to transformers - maybe that can help ( https://github.com/vasqu/mamba2-torch but look into the open PR if you want correct batched inference).

Got it, thanks again for the explanations (and enjoy the holidays)!

Any reason to withhold this PR? Seems like Bamba (#34982) has been merged and silently merged this PR's fixes (https://github.com/huggingface/transformers/pull/34982/commits/856cb3a41f9131f9f1c35fbe250dfdd2f38a8368) without a (commit) reference / dubious credit (i.e. adopted this "refactor" but fixed bugs themself?).

Failing tests are unrelated (internal hub failures).

@molbap @ArthurZucker 

Hi @vasqu, this should be merged soon! In terms of precedence I'll add an `From PR ... by @vasqu` in comment so future users understand where the fixes come from, not a worry
ping @ArthurZucker  for merge! 

Hey 👋 I don't need direct credit, I just think that the list given in the docstring is misleading: 
https://github.com/huggingface/transformers/blob/667ed5635e6fd7e2df4fc23012746b1c0cbb7575/src/transformers/models/bamba/modular_bamba.py#L214-L218
The changes are mainly because of the cache + dropping some attributes.