Possible reshape error in Mamba2Mixer causing inference issue
Author: Kirire
Pull request: no
Labels: name: bug

### System Info

Hello,  
I'm encountering a dimension mismatch error when running inference with the `Mamba2ForCausalLM` class. The error originates from line 579 of the `Mamba2Mixer` class on the `main` branch:

```python
D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size)
```

Could it be that there's a dimension mix-up a few lines above (line 572)? Specifically, this line:

```python
hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float()
```

Shouldn't it be:

```python
hidden_states = hidden_states.reshape(batch_size, seq_len, self.head_dim, -1).float()
```

Thanks in advance,  
Cyrile

### Who can help?

_No response_

### Information

- [x] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

```python
import torch
from transformers import Mamba2Config, Mamba2ForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("TempestTeam/TempestLLM")
config = Mamba2Config()
config.vocab_size = len(tokenizer)
config.bos_token_id = tokenizer.encode(tokenizer.special_tokens_map['bos_token'])[0]
config.eos_token_id = tokenizer.encode(tokenizer.special_tokens_map['eos_token'])[0]
config.pad_token_id = tokenizer.encode(tokenizer.special_tokens_map['pad_token'])[0]
config.num_hidden_layers = 1
config.hidden_size = 16
config.num_heads = 4
config.head_dim = 4
config.state_size = 8

model = Mamba2ForCausalLM(config)

tok = tokenizer("Hello World!", return_tensors="pt")

output = model(**tok)
```

### Expected behavior

An inference pass that completes without any dimension mismatch errors.

Comments:
Nope, not a mistake. I'll explain in a second but here would be better default values you could use:
```python
from transformers import Mamba2Config, Mamba2ForCausalLM, AutoTokenizer


tokenizer = AutoTokenizer.from_pretrained("TempestTeam/TempestLLM")
config = Mamba2Config()
config.vocab_size = len(tokenizer)
config.bos_token_id = tokenizer.encode(tokenizer.special_tokens_map['bos_token'])[0]
config.eos_token_id = tokenizer.encode(tokenizer.special_tokens_map['eos_token'])[0]
config.pad_token_id = tokenizer.encode(tokenizer.special_tokens_map['pad_token'])[0]
config.num_hidden_layers = 1
config.hidden_size = 8
config.num_heads = 4
config.head_dim = 4
config.state_size = 8
config.n_groups = 1

model = Mamba2ForCausalLM(config)

tok = tokenizer("Hello World!", return_tensors="pt")

output = model(**tok)

```

Why does it fail in the first place? Because your current config is not aligning with what's allowed in the mamba2 dimensions; in this specific case, it's because of `hidden_size * expand != num_heads * head_dim` or with the values inserted `16 * 2 (32) != 4 * 4 (16)`. My best guess is that you weren't aware of the expand factor ;) 

I must admit tho that it's hard to know which conditions might cause issues without knowing the architecture (by heart). 

Ok! Thank you. It's work 👍 
You're right — in a Transformer structure, it's usually `num_heads * dim_head = hidden_dim`. Hence my mistake 😅  
Maybe adding a dimension check in the `Mamba2Config` class could help guide future users of this architecture?

Also, unlike Transformers where the attention mechanism operates in the original hidden dimension and only the FeedForward part expands it, Mamba2 actually "works" in the expanded dimension directly. That inversion can definitely be a source of confusion at first.

I think adding a check on creation time doesn't hurt and might help users coming from a different view like you said. We usually raise a value error in such cases then.

Would you be willing to open a PR for that? 

cc @molbap 

Alright, I'm on it!

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Was fixed by #37589 so closing it