[`Add Mamba`] Adds support for the `Mamba` models
Author: ArthurZucker
Pull request: yes

# What does this PR do?
- [x] Implement cpu ops 
- [x] Add integration tests 
- [x] Implement fast path 
- [x] check training + peft 
- [x] convert all checkpoints: just need to make sure config is correct

Feel free to try this: 
```python 
from transformers import MambaConfig, MambaForCausalLM, AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("ArthurZ/mamba-130m")
tokenizer.pad_token = tokenizer.eos_token

model = MambaForCausalLM.from_pretrained("ArthurZ/mamba-130m", vocab_size=50280, num_hidden_layers=24, torch_dtype=torch.float32)
model.config.use_cache = True
input_ids = tokenizer("Hey how are you doing?", return_tensors= "pt")["input_ids"]

out = model.generate(input_ids, max_new_tokens=10)
print(tokenizer.batch_decode(out))
```

Peft training that works, thanks @younesbelkada : Results: https://huggingface.co/ArthurZ/mamba-2.4b-english-quotes 
```python 
from datasets import load_dataset
from trl import SFTTrainer
from peft import LoraConfig
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
model_id = "ArthurZ/mamba-2.8b"
tokenizer = AutoTokenizer.from_pretrained(model_id, pad_token ="
")
model = AutoModelForCausalLM.from_pretrained(model_id)
dataset = load_dataset("Abirate/english_quotes", split="train")
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    logging_dir='./logs',
    logging_steps=10,
    learning_rate=2e-3
)
lora_config =  LoraConfig(
        r=8,
        target_modules="all-linear",
        task_type="CAUSAL_LM",
        bias="none"
)
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    peft_config=lora_config,
    train_dataset=dataset,
    dataset_text_field="quote",
)
trainer.train()
```

pink: 360m, full fine-tune
bleue:  2.8b peft 
red: 2.8b peft
<img width="480" alt="image" src="https://github.com/huggingface/transformers/assets/48595927/9e5f5726-eb39-4aaf-859e-2f9b8a02ffb3">


fixes #28086

Comments:
Oups! Still planned but KVCache will come first

Alright I am picking this back up! 

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

Hey, it's great to see that mamba is being integrated in Transformers! Just wondering, is there a timeline or ETA for this PR? Thanks so much.

I want to merge it asap so probably max end of next week! 

Got side tracked, done with caching issues! 
Was meditating the stateful vs stateless approach we want to take to support torch compile and graphs without the extra complexity similarly to #27931. 
It was advised that for mamba, cache should work in a stateless manner


Hello, is this one fine and have same generate speed as the `mamba/mamba_ssm/utils/generation.py `. Someone like mamba-hf used mini-mamba which has a bad speed when training and genrateing. And the original mamba generate code will lead to a nan error after the last token is pad. But i see your code `tokenizer.pad_token = tokenizer.eos_token` here. Is it necessary? In my test, i set the different token, sometimes will lead the error as i mentioned.

Hey! API will change today / tomorrow to a statefull API, with setup cache instead of moving around the inference params. 
Will support both and benchmark both. 
Could you share a minimal reproducer of your code to run? 

> Hey! API will change today / tomorrow to a statefull API, with setup cache instead of moving around the inference params. Will support both and benchmark both. Could you share a minimal reproducer of your code to run?

here is the code. i copy modeling_mamba.py and configuration_mamba.py  from your branch.

```python
from modeling_mamba import MambaForCausalLM
from configuration_mamba import MambaConfig


config = MambaConfig(vocab_size=256,
        hidden_size=768,
        state_size=16,
        num_hidden_layers=32,
        layer_norm_epsilon=1e-5,
        pad_token_id=0,
        bos_token_id=1,
        eos_token_id=2,
        expand=2,
        dt_rank="auto",
        tie_word_embeddings=True,
        use_bias=False,
        use_conv_bias=True,
        hidden_act="silu",
        initializer_range=0.1,
        residual_in_fp32=True,
        use_cache=True,)

model = MambaForCausalLM(config).to("cuda")

import torch

input_ids = torch.tensor([[1, 2, 3, 4, 5]]).to("cuda")

model(input_ids=input_ids,
        inputs_embeds = None,
        inference_params = None,
        labels = None,
        output_hidden_states = None,
        return_dict = None,
        )

```

Alright just wait a bit the PR is still draft :) 

I'll finish it today / tomorrow! 

😅 sorry all! Finally back on track!
ETA: merged tomorrow 🥳 

A small update here, latest push supports peft training

Todo:

- [x] gradient checkpointing CIs
- [ ] finish few remaining CIs + PR documentation
- [ ] init scheme
- [x] properly handle slow and fast path: a single class is enough
- [ ] torch compile benchmarks 


> Todo:
> 
> * gradient checkpointing CIs
> * properly handle CPU init / force move to cuda
> * torch compile benchmarks
> * profile for potential memory leaks
> * finish few remaining CIs + PR documentation

sorry that when i use the code to generate, it will get error
```
out = model.generate(input_ids=test_input_ids["input_ids"].to("cuda"),max_length=128,pad_token_id=tokenizer.pad_token_id)
print(tokenizer.decode(out[0],skip_special_tokens=True))
```

```
Unexpected exception formatting exception. Falling back to standard exception
Traceback (most recent call last):
  File "<string>", line 21, in _selective_scan_update_kernel
KeyError: ('2-.-0-.-0-83ca8b715a9dc5f32dc1110973485f64-d6252949da17ceb5f3a278a70250af13-1af5134066c618146d2cd009138944a0-14de7de5c4da5794c8ca14e7e41a122d-3498c340fd4b6ee7805fd54b882a04f5-e1f133f98d04093da2078dfc51c36b72-b26258bf01f839199e39d64851821f26-d7c06e3b46e708006c15224aac7a1378-f585402118c8a136948ce0a49cfe122c', (torch.float32, torch.float32, torch.float32, torch.float32, torch.float32, torch.float32, torch.float32, torch.float32, torch.float32, torch.float32, 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32', 'i32'), (True, 32, True, True, True, 16), (True, True, True, True, True, True, True, True, True, True, (False, True), (True, False), (True, False), (True, False), (True, False), (False, True), (True, False), (False, True), (True, False), (False, True), (False, True), (True, False), (False, True), (False, False), (False, True), (False, False), (False, True), (False, True), (True, False), (False, True), (True, False), (False, True)))

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 937, in build_triton_ir
    generator.visit(fn.parse())
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
    return super().visit(node)
  File "/root/miniconda3/lib/python3.8/ast.py", line 371, in visit
    return visitor(node)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 183, in visit_Module
    ast.NodeVisitor.generic_visit(self, node)
  File "/root/miniconda3/lib/python3.8/ast.py", line 379, in generic_visit
    self.visit(item)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
    return super().visit(node)
  File "/root/miniconda3/lib/python3.8/ast.py", line 371, in visit
    return visitor(node)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 252, in visit_FunctionDef
    has_ret = self.visit_compound_statement(node.body)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 177, in visit_compound_statement
    self.last_ret_type = self.visit(stmt)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
    return super().visit(node)
  File "/root/miniconda3/lib/python3.8/ast.py", line 371, in visit
    return visitor(node)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 486, in visit_If
    self.visit_compound_statement(node.body)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 177, in visit_compound_statement
    self.last_ret_type = self.visit(stmt)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
    return super().visit(node)
  File "/root/miniconda3/lib/python3.8/ast.py", line 371, in visit
    return visitor(node)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 301, in visit_Assign
    values = self.visit(node.value)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
    return super().visit(node)
  File "/root/miniconda3/lib/python3.8/ast.py", line 371, in visit
    return visitor(node)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 757, in visit_Call
    args = [self.visit(arg) for arg in node.args]
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 757, in <listcomp>
    args = [self.visit(arg) for arg in node.args]
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
    return super().visit(node)
  File "/root/miniconda3/lib/python3.8/ast.py", line 371, in visit
    return visitor(node)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 751, in visit_Call
    fn = self.visit(node.func)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
    return super().visit(node)
  File "/root/miniconda3/lib/python3.8/ast.py", line 371, in visit
    return visitor(node)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 835, in visit_Attribute
    lhs = self.visit(node.value)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
    return super().visit(node)
  File "/root/miniconda3/lib/python3.8/ast.py", line 371, in visit
    return visitor(node)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 839, in visit_Attribute
    return getattr(lhs, node.attr)
AttributeError: module 'triton.language' has no attribute 'math'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/root/miniconda3/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 3460, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "/tmp/ipykernel_178048/3165352146.py", line 3, in <module>
    out = model.generate(input_ids=test_input_ids["input_ids"].to("cuda"),max_length=128,pad_token_id=tokenizer.pad_token_id)
  File "/root/miniconda3/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/miniconda3/lib/python3.8/site-packages/transformers/generation/utils.py", line 1544, in generate
    return self.greedy_search(
  File "/root/miniconda3/lib/python3.8/site-packages/transformers/generation/utils.py", line 2404, in greedy_search
    outputs = self(
  File "/root/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/test-mamba/modeling_mamba.py", line 650, in forward
    mamba_outputs = self.backbone(
  File "/root/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/test-mamba/modeling_mamba.py", line 548, in forward
    hidden_states = mixer_block(hidden_states, inference_params=inference_params)
  File "/root/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/test-mamba/modeling_mamba.py", line 311, in forward
    hidden_states = self.mixer(hidden_states, inference_params=inference_params)
  File "/root/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/test-mamba/modeling_mamba.py", line 152, in forward
    y = selective_state_update(
  File "/root/miniconda3/lib/python3.8/site-packages/mamba_ssm/ops/triton/selective_state_update.py", line 137, in selective_state_update
    _selective_scan_update_kernel[grid](
  File "/root/miniconda3/lib/python3.8/site-packages/triton/runtime/autotuner.py", line 199, in run
    return self.fn.run(*args, **kwargs)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/runtime/autotuner.py", line 199, in run
    return self.fn.run(*args, **kwargs)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/runtime/autotuner.py", line 199, in run
    return self.fn.run(*args, **kwargs)
  [Previous line repeated 1 more time]
  File "<string>", line 41, in _selective_scan_update_kernel
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 1620, in compile
    next_module = compile(module)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 1549, in <lambda>
    lambda src: ast_to_ttir(src, signature, configs[0], constants)),
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 962, in ast_to_ttir
    mod, _ = build_triton_ir(fn, signature, specialization, constants)
  File "/root/miniconda3/lib/python3.8/site-packages/triton/compiler.py", line 942, in build_triton_ir
    raise CompilationError(fn.src, node) from e
triton.compiler.CompilationError: at 58:34:
def _selective_scan_update_kernel(
    # Pointers to matrices
    state_ptr, x_ptr, dt_ptr, dt_bias_ptr, A_ptr, B_ptr, C_ptr, D_ptr, z_ptr, out_ptr,
    # Matrix dimensions
    batch, dim, dstate,
    # Strides
    stride_state_batch, stride_state_dim, stride_state_dstate,
    stride_x_batch, stride_x_dim,
    stride_dt_batch, stride_dt_dim,
    stride_dt_bias_dim,
    stride_A_dim, stride_A_dstate,
    stride_B_batch, stride_B_dstate,
    stride_C_batch, stride_C_dstate,
    stride_D_dim,
    stride_z_batch, stride_z_dim,
    stride_out_batch, stride_out_dim,
    # Meta-parameters
    DT_SOFTPLUS: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    HAS_DT_BIAS: tl.constexpr,
    HAS_D: tl.constexpr,
    HAS_Z: tl.constexpr,
    BLOCK_SIZE_DSTATE: tl.constexpr,
):
    pid_m = tl.program_id(axis=0)
    pid_b = tl.program_id(axis=1)
    state_ptr += pid_b * stride_state_batch
    x_ptr += pid_b * stride_x_batch
    dt_ptr += pid_b * stride_dt_batch
    B_ptr += pid_b * stride_B_batch
    C_ptr += pid_b * stride_C_batch
    if HAS_Z:
        z_ptr += pid_b * stride_z_batch
    out_ptr += pid_b * stride_out_batch

    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = tl.arange(0, BLOCK_SIZE_DSTATE)
    state_ptrs = state_ptr + (offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate)
    x_ptrs = x_ptr + offs_m * stride_x_dim
    dt_ptrs = dt_ptr + offs_m * stride_dt_dim
    if HAS_DT_BIAS:
        dt_bias_ptrs = dt_bias_ptr + offs_m * stride_dt_bias_dim
    A_ptrs = A_ptr + (offs_m[:, None] * stride_A_dim + offs_n[None, :] * stride_A_dstate)
    B_ptrs = B_ptr + offs_n * stride_B_dstate
    C_ptrs = C_ptr + offs_n * stride_C_dstate
    if HAS_D:
        D_ptrs = D_ptr + offs_m * stride_D_dim
    if HAS_Z:
        z_ptrs = z_ptr + offs_m * stride_z_dim
    out_ptrs = out_ptr + offs_m * stride_out_dim

    state = tl.load(state_ptrs, mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate), other=0.0)
    x = tl.load(x_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
    dt = tl.load(dt_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
    if HAS_DT_BIAS:
        dt += tl.load(dt_bias_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
    if DT_SOFTPLUS:
        dt = tl.where(dt <= 20.0, tl.math.log1p(tl.exp(dt)), dt)
                                  ^

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/root/miniconda3/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 2057, in showtraceback
    stb = self.InteractiveTB.structured_traceback(
  File "/root/miniconda3/lib/python3.8/site-packages/IPython/core/ultratb.py", line 1288, in structured_traceback
    return FormattedTB.structured_traceback(
  File "/root/miniconda3/lib/python3.8/site-packages/IPython/core/ultratb.py", line 1177, in structured_traceback
    return VerboseTB.structured_traceback(
  File "/root/miniconda3/lib/python3.8/site-packages/IPython/core/ultratb.py", line 1049, in structured_traceback
    formatted_exceptions += self.format_exception_as_a_whole(etype, evalue, etb, lines_of_context,
  File "/root/miniconda3/lib/python3.8/site-packages/IPython/core/ultratb.py", line 935, in format_exception_as_a_whole
    self.get_records(etb, number_of_lines_of_context, tb_offset) if etb else []
  File "/root/miniconda3/lib/python3.8/site-packages/IPython/core/ultratb.py", line 1003, in get_records
    lines, first = inspect.getsourcelines(etb.tb_frame)
  File "/root/miniconda3/lib/python3.8/inspect.py", line 979, in getsourcelines
    lines, lnum = findsource(object)
  File "/root/miniconda3/lib/python3.8/inspect.py", line 798, in findsource
    raise OSError('could not get source code')
OSError: could not get source code
```

![image](https://github.com/huggingface/transformers/assets/53668275/5b83d41b-e7e5-4c47-b83e-12de3e1012d4)



I am getting: `Hello my name is Dâ€™Amnonâ€™,â€� said the voice of the dâ€™gone. â€œI am the dâ€™giver of lifeâ€™ on Earthâ€� said the dâ€™giver of life. â€œI am the dâ€™giver of lifeâ€™ on earthâ€� said the dâ€™giver of life. â€œI am the dâ€™giver of lifeâ€™ on earthâ€� said the dâ€™giver of life. â€œI am the dâ€™giver of` 


> Todo:
> 
> * gradient checkpointing CIs
> * properly handle CPU init / force move to cuda
> * torch compile benchmarks
> * profile for potential memory leaks
> * finish few remaining CIs + PR documentation

Maybe need to remove the InferenceParams. Merge it to the MambaConfig. 

No that's not really the way we handle the `cache`. I will most probably rename it to `cache`

> I am getting: `Hello my name is Dâ€™Amnonâ€™,â€� said the voice of the dâ€™gone. â€œI am the dâ€™giver of lifeâ€™ on Earthâ€� said the dâ€™giver of life. â€œI am the dâ€™giver of lifeâ€™ on earthâ€� said the dâ€™giver of life. â€œI am the dâ€™giver of lifeâ€™ on earthâ€� said the dâ€™giver of life. â€œI am the dâ€™giver of`

in https://github.com/huggingface/transformers/blob/edb4e91a7c34b961adb412c52ed97740b79dc5d7/src/transformers/models/mamba/modeling_mamba.py#L602

KeyError
'inference_params'
  File "/root/test-mamba/modeling_mamba.py", line 608, in _update_model_kwargs_for_generation
    # if model_kwargs not in outputs, we need to update it
  File "/tmp/ipykernel_191194/1978564718.py", line 3, in <module>
    output = model.generate(test_input_ids, max_length=512, num_return_sequences=1,pad_token_id=0)
KeyError: 'inference_params'

![image](https://github.com/huggingface/transformers/assets/53668275/dbc83797-a042-40d1-92d8-ee43cbfe4cf6)


the code is 
![image](https://github.com/huggingface/transformers/assets/53668275/db116998-6f01-4c66-903c-2d7d54c5d731)

@relic-yuexi do you have discord or somewhere else you can report this? This way we don't burden the PR with our tests 🤗 

Alright slow version matches fast version. For now I think this looks good. 
TODO left for this PR: 
- [x] test inner fn 
- [x] final CIs 

TODO in a follow up PR, very interesting work mostly for the kernel-less implementation:
 
- [ ] custom backward: save for backward to explicitly re-compute instead of loading big big tensors (slow forward)
- [ ] use parallel scan for slow forward (Blelloch parallel scan algorithm)
- [ ] use recomputation-trick in the Blelloch parallel scan algorithm

Done! 🤗 

@ArthurZucker Thank you for this amazing addition. Are there any plans to add something equivalent to `attention_mask` for Mamba?

not sure why would you need it?

- For batched inference with inputs of different length.
- For pretraining with different masking schemes than a causal mask. 

There is no notion of causal mask or masking in `mamba` as it is not based on attention. That's why I am not sure I follow

Hi. 

There is a problem in the Trainer where the `logits` returned by `Trainer.prediction_step` will return a `tuple[Tensor, MambaCache]` object. This causes a host of issues when accelerate tries to move the logits on the same device, change datatypes, etc. The solution is to set the "keys_to_ignore_at_inference" field of the associated Config class to include "cache_params". The change is simple:

```
class MambaConfig:
    keys_to_ignore_at_inference = ["cache_params"]
```

Full disclosure, I encountered this "bug" in my own `MambaForSequenceClassification` class, not a module from transformers itself and I have not really tested this thoroughly to see if it is present in the classes from transformers.

@ArthurZucker tagging you :)

Feel free to open a PR for the fix! 🤗 

Also `use_cache=False` should prevent this as well no?