[Add Mamba] Adds support for the Mamba models

[ArthurZucker]

What does this PR do?

• Implement cpu ops
• Add integration tests
• Implement fast path
• check training + peft
• convert all checkpoints: just need to make sure config is correct

Feel free to try this:

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

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 ="<s>")
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

fixes #28086

[ArthurZucker]

Oups! Still planned but KVCache will come first

[ArthurZucker]

Alright I am picking this back up!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[apoorvkh]

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.

[ArthurZucker]

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

[ArthurZucker]

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

[ArthurZucker]

Done! 🤗

[LysandreJik]

This looks good to me, please add the example that you have in the PR description somewhere in the documentation as well. The current examples don't really show how to use the model imo.

[abdulfatir]

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

[ArthurZucker]

not sure why would you need it?

[abdulfatir]

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

[ArthurZucker]

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

[lkurlandski]

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 :)

[ArthurZucker]

Feel free to open a PR for the fix! 🤗

[ArthurZucker]

Also use_cache=False should prevent this as well no?