Add config validation and style tweaks
Author: Kirire
Pull request: yes

**Summary of changes:**

1. **Renamed `input_states` to `hidden_states`** in the `torch_forward` method signature and body, to align with HuggingFace Transformers' naming conventions:
   ```python
   def torch_forward(self, hidden_states: torch.Tensor, ...)
   ```
   This change improves consistency with the rest of the codebase and aligns better with the standard used in other HuggingFace models.

2. **Removed redundant attention masking** from the `forward` method:
   ```python
   if attention_mask is not None ...
       hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
   ```
   This operation is already handled later in `torch_forward`, so it's no longer necessary here. Removing it avoids potential duplication and keeps the logic centralized.

3. **Added a configuration validation check in `Mamba2Config.__init__`**:
   ```python
   if hidden_size * expand != num_heads * head_dim:
       raise AttributeError(...)
   ```
   This ensures that the configuration is internally consistent, and helps users catch misconfigurations early by raising a clear error during initialization.

Related to #37554

4. **Improved Loss Handling During Gradient Accumulation**

Modifies the loss computation by replacing the hardcoded use of `CrossEntropyLoss`:

```python
loss_fct = CrossEntropyLoss()
loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
```

with a more flexible call to a configurable `loss_function`:

```python
loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
```

This change improves support for gradient accumulation scenarios and allows for easier customization of model-specific loss functions.

Related to #34191

Comments:
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the `Ready for review` button (at the bottom of the PR page). This will assign reviewers and trigger CI.

cc @vasqu @molbap from #37554

The failing CI can be fixed with `make style` (after having the deps installed e.g. `pip install -e ".[quality]"`).

Hi! I ran make quality, everything should be good now 👍

Just a last thing: Could you add some test to make sure it works as intended? #33316 is a good reference for how this can be done.

Gentle ping @Kirire 

@Kirire are you still interested in finishing this PR? 

Hi! Sorry for the late reply — I recently became a dad, so things have been a bit busy on my end 👶🙂
But yes, I'm definitely still interested, and I’ll get back to it starting tomorrow!

First of all, congrats to you and your family !!! :partying_face: And no worries, take your time.

Thanks a lot! 😊
I've now added the test for the configuration — let me know if that works for you or if you'd like any changes!

oups my bad ^^'

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

@Kirire thanks for contributing! And enjoy the days with your family :wink: