Batch Generation giving different output when using batch size > 1 or when using padding in MambaForCausalLM

[piyushdevlpr]

System Info

• transformers version: 4.41.2
• Platform: Linux-6.5.0-41-generic-x86_64-with-glibc2.35
• Python version: 3.10.8
• Huggingface_hub version: 0.23.3
• Safetensors version: 0.4.3
• Accelerate version: 0.26.0
• Accelerate config: not found
• PyTorch version (GPU?): 2.0.0+cu117 (True)
• Tensorflow version (GPU?): 2.9.0 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using GPU in script?: Yes
• Using distributed or parallel set-up in script?: No

Who can help?

@ArthurZucker @gante

Information

• [ ] The official example scripts
• [X] 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

I have trained a MambaForCausalLM model on a custom dataset. I am using the following code to generate next token in eval mode -

model = MambaForCausalLM.from_pretrained("mamba_custom")
tokenizer = PreTrainedTokenizerFast.from_pretrained("mamba_tokenizer_custom")
tokenizer.padding_side = 'left'

inputs_rl = tokenizer(sentence, padding="max_length", truncation=True, max_length=100, return_tensors="pt").to("cuda:0")
with torch.no_grad():
    outputs = model(inputs_rl["input_ids"], attention_mask=inputs_rl["attention_mask"])

Expected behavior

The tokenizer pads the input towards the left side. When I change the argument padding="max_length" to generate inputs without padding, I get different tokens as prediction.

Using model.generate gives the same issue as well.

[amyeroberts]

cc @gante

[gante]

@piyushdevlpr 👋

Mamba, contrarily to transformers models, does not take an attention mask as input (see the signature here). As such, it does not support padding, and will return different values.

(I'm going to open a PR to try to prevent this issue from happening again)