RecurrentGemma Doesn't Support left padding?
Author: godjw
Pull request: no
Labels: name: Text

### System Info

- `transformers` version: 4.40.0
- Python version: 3.9.9
- Huggingface_hub version: 0.23.0
- Safetensors version: 0.4.3
- Accelerate version: 0.29.3
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.2.2+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed

### Who can help?

@ArthurZucker @younesbelkada 

### 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 am currently trying to finetune the RecurrentGemmaModel.

However, when I train the model with left padding the training outputs where quite strange so I tried to debug it.

I checked this part about the causal mask.
https://github.com/huggingface/transformers/blob/96eb06286b63c9c93334d507e632c175d6ba8b28/src/transformers/models/recurrent_gemma/modeling_recurrent_gemma.py#L753-L779


When I put a left padded dummy attention mask like below, the attention mask looks strange.


```python
dummy_attention_mask = tensor([[0, 0, 1, 1],
        [1, 1, 1, 1]], device='cuda:0')

inputs_embeds = model.model.embed_tokens(input_ids.to("cuda"))
hidden_states = inputs_embeds
cache_position = torch.arange(hidden_states.shape[1])

mask_output = model.model._update_causal_mask(dummy_attention_mask, attention_mask.to('cuda'), inputs_embeds, cache_position.to("cuda"))
print(mask_output)
tensor([[[[-0.0000e+00, -0.0000e+00, -0.0000e+00, -0.0000e+00],
          [-0.0000e+00, -0.0000e+00, -0.0000e+00, -0.0000e+00],
          [-3.4028e+38, -3.4028e+38, -0.0000e+00, -3.4028e+38],
          [-3.4028e+38, -3.4028e+38, -0.0000e+00, -0.0000e+00]]],


        [[[-0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38],
          [-0.0000e+00, -0.0000e+00, -3.4028e+38, -3.4028e+38],
          [ 0.0000e+00, -0.0000e+00, -0.0000e+00, -3.4028e+38],
          [ 0.0000e+00,  0.0000e+00, -0.0000e+00, -0.0000e+00]]]],
       device='cuda:0')
```


However, giving the right padded attention mask produces the right attention mask. 
```python
dummy_attention_mask = tensor([[ 1,  1, 0, 0],
        [ 1, 1, 1, 1]], device='cuda:0')

# output
tensor([[[[-0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38],
          [-0.0000e+00, -0.0000e+00, -3.4028e+38, -3.4028e+38],
          [ 0.0000e+00, -0.0000e+00, -3.4028e+38, -3.4028e+38],
          [ 0.0000e+00,  0.0000e+00, -3.4028e+38, -3.4028e+38]]],


        [[[-0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38],
          [-0.0000e+00, -0.0000e+00, -3.4028e+38, -3.4028e+38],
          [ 0.0000e+00, -0.0000e+00, -0.0000e+00, -3.4028e+38],
          [ 0.0000e+00,  0.0000e+00, -0.0000e+00, -0.0000e+00]]]],
```


### Expected behavior

I think The expected first mask which is left padded should look like below, because the padded parts should not be attended.

```python 
# Expected Mask
[[[[-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
          [-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],
          [-3.4028e+38, -3.4028e+38, -0.0000e+00, -3.4028e+38],
          [-3.4028e+38, -3.4028e+38, -0.0000e+00, -0.0000e+00]]],
```

Comments:
Hi @godjw and thanks for the issue, indeed, it seems -at least from my local setup- that left-padding has issues in `float16` precision, does that match your finetuning setup?
The issue does seem to arise since v4.40 where we introduced RecurrentGemma, despite having a test in modeling that does check for left-padded generation. ccing @ArthurZucker for vis and will take a look tomorrow!

@molbap 
Thanks for the reply! and I am using the `float32` precision for my finetuning setup.
But I think the precisions will not make much difference for the padding mask. 

Currently the finetuning only works with `right padding` 

This seems expected in terms of mask creation:
- if you are using SDPA and left padding, the first 2 rows are going to be `-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38`, but this is not supported in sdpa so we switch to full attention,  `-0.0000e+00, -0.0000e+00, -0.0000e+00, -0.0000e+00` which should not affect the result
- however, if you are using `output_attention=True` when training / not using `sdpa` in the `attn_implementation` then this is not expected. 

Now the causal mask is not going to be an issue TBH, but the way the model works might. Since there is a convolution layer in the RNN part, we might need to make it ignore padding. 
One more think is which padding token is used, and if the embedding was resized for it, the embedding value needs to be an average of all the model's embeddings

@ArthurZucker 
I'm sorry, but I'm having some difficulty understanding a few points. Could you please clarify why the first two rows becoming `-3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38` would not work with `sdpa` and why switching to full attention wouldn't affect the result? It seems quite unusual to me that the padding tokens would be attended to. 

I was able to train and inference with the model by giving the model no attention mask and just using the default causal mask, which doesn't care about the padding tokens- but with the left padded attention mask only training was possible and the inference results were really bad.

Additionally, I wanted to mention that I am using `sdpa` because it is the default and only supported type of attention in the `RecurrentGemmaModel`.

Thank you very much for your help!



This piece of code in `modeling_llama.py` should help you:
```
# Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
            # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
            # Details: https://github.com/pytorch/pytorch/issues/110213
            causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
```

🤗 

Thank you very much!
Then the problem could be because the padding tokens are not considered properly, as you've mentioned.

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.