Remove ambiguous `padding_mask` and instead use a 2D->4D Attn Mask Mapper
Author: patrickvonplaten
Pull request: yes

# What does this PR do?

For models that have Flash Attention 2 (FA2) implemented we currently pass both `padding_mask` and `attention_mask` to the respective vanilla attention class, *e.g.* `LlamaAttention` and to the FA2 class, *e.g.* `LlamaFlashAttention2`.

**However**, `padding_mask` is **not** used for `LlamaAttention` and `attention_mask` is not used for `LlamaFlashAttention2`. Conceptually the two masks are the same, only that `attention_mask` is a 4D mask while `padding_mask` is a 2D mask.

Passing around both masks and having both masks as concepts in our codebase is ambiguous and hurts readability. In this PR, I propose to remove the concept of `padding_mask` completely and instead just pass either a 2D or 4D `attention_mask` depending on whether we use FA2 or not.

**Note**: An additional benefit of this PR is that it will improve the performance when using FA2 as we will not create a 4D attention mask anymore.

**Benchmarks**:

The following script was used to benchmark the effect this mask implementation has on forward and generate.

```py
#!/usr/bin/env python3
from transformers import AutoTokenizer, AutoModelForCausalLM
import time
import torch

DEVICE = "cuda:1"

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf", torch_dtype=torch.float16, low_cpu_mem_usage=True)
model.to(DEVICE)


# forward
print("Forward benchmarks")
print(50 * "=")

for batch_size in (1, 4, 16):
    for input_seq in (4, 16, 256):
        input_ids = torch.ones((batch_size, input_seq), dtype=torch.long, device=DEVICE)
        attention_mask = torch.ones_like(input_ids)
        attention_mask[0, 3] = 0

        times = []
        for _ in range(3):
            start_time = time.time()
            with torch.no_grad():
                logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
            times.append(time.time() - start_time)

        result = min(times)

        print(f"Forward bsz={batch_size}, input_seq={input_seq}: {result}")


# generate
print("Generate benchmarks")
print(50 * "=")

for batch_size in (1, 16):
    for input_seq in (4, 256):
        input_ids = torch.ones((batch_size, input_seq), dtype=torch.long, device=DEVICE)
        attention_mask = torch.ones_like(input_ids)
        attention_mask[0, 3] = 0

        times = []
        for _ in range(3):
            start_time = time.time()
            out = model.generate(input_ids=input_ids, max_new_tokens=256, do_sample=False)
            times.append(time.time() - start_time)

        result = min(times)

        print(f"Generate bsz={batch_size}, input_seq={input_seq}: {result}")
```

**This PR:**
```
Forward benchmarks
==================================================
Forward bsz=1, input_seq=4: 0.012479066848754883
Forward bsz=1, input_seq=16: 0.011297464370727539
Forward bsz=1, input_seq=256: 0.01240849494934082
Forward bsz=4, input_seq=4: 0.011190414428710938
Forward bsz=4, input_seq=16: 0.013025283813476562
Forward bsz=4, input_seq=256: 0.03526663780212402
Forward bsz=16, input_seq=4: 0.01126551628112793
Forward bsz=16, input_seq=16: 0.012389421463012695
Forward bsz=16, input_seq=256: 0.1560053825378418
Generate benchmarks
==================================================
Generate bsz=1, input_seq=4: 4.527426719665527
Generate bsz=1, input_seq=256: 4.667049169540405
Generate bsz=16, input_seq=4: 5.524803400039673
Generate bsz=16, input_seq=256: 7.931211709976196
```

**Current main:**
```
Forward benchmarks
==================================================
Forward bsz=1, input_seq=4: 0.017528295516967773
Forward bsz=1, input_seq=16: 0.012105464935302734
Forward bsz=1, input_seq=256: 0.01315617561340332
Forward bsz=4, input_seq=4: 0.011912107467651367
Forward bsz=4, input_seq=16: 0.013910531997680664
Forward bsz=4, input_seq=256: 0.035504817962646484
Forward bsz=16, input_seq=4: 0.012083053588867188
Forward bsz=16, input_seq=16: 0.012537956237792969
Forward bsz=16, input_seq=256: 0.15653300285339355
Generate benchmarks
==================================================
Generate bsz=1, input_seq=4: 4.554980516433716
Generate bsz=1, input_seq=256: 4.695344686508179
Generate bsz=16, input_seq=4: 5.55778431892395
Generate bsz=16, input_seq=256: 7.969247102737427
```

=> We don't see any drop in performance at all.

I've verified that the following tests all pass on a single GPU (RTX4090):
FA2:
```
RUN_SLOW=1 pytest -m flash_attn_test tests/models/llama/test_modeling_llama.py
```
and all Llama fast tests:
```
CUDA_VISIBLE_DEVICES="0" RUN_SLOW=1 pytest tests/models/llama/test_modeling_llama.py
```

Comments:
_The documentation is not available anymore as the PR was closed or merged._

> Thank you! I am happy with it, just wondering whether changing the `attention_mask` input from being 4D to 2D in `LlamaDecoderLayer` & `LlamaAttention` is considered a breaking change or not.

To me they are internal classes and thus changing the format of the attention_mask is ok. @LysandreJik @younesbelkada @ArthurZucker what do you think? 




This PR should help make the following PRs nicer / cleaner:
- https://github.com/huggingface/transformers/pull/26722
- https://github.com/huggingface/transformers/pull/26572

> LGTM, we might have to keep some logic to pop the padding mask for 1 release for BC. let's do a deprecation cycle no?

They are all internal classes so BC doesn't really apply here IMO. Also it's only been in one release, so not sure we need to have a deprecation cycle here

Ok, it's true that it has only been 2 weeks, good for me


@patrickvonplaten I disagree that the non-top-level classes are internal. They are publicly exposed (no `_`), and have docstrings. There are a few libraries that were broken by the addition of padding_mask here: https://github.com/huggingface/transformers/pull/25598

e.g.
https://github.com/casper-hansen/AutoAWQ/pull/88
https://github.com/mosaicml/llm-foundry/pull/643
https://github.com/hpcaitech/ColossalAI/pull/4908

so IMO modifying the non-top level signature / expected args (or args format, e.g. changing an arg from 4D to 2D) classes should be done with great care.

> @patrickvonplaten I disagree that the non-top-level classes are internal. They are publicly exposed (no `_`), and have docstrings. There are a few libraries that were broken by the addition of padding_mask here: #25598
> 
> e.g. [casper-hansen/AutoAWQ#88](https://github.com/casper-hansen/AutoAWQ/pull/88) [mosaicml/llm-foundry#643](https://github.com/mosaicml/llm-foundry/pull/643) [hpcaitech/ColossalAI#4908](https://github.com/hpcaitech/ColossalAI/pull/4908)
> 
> so IMO modifying the non-top level signature / expected args (or args format, e.g. changing an arg from 4D to 2D) classes should be done with great care.

This PRs all kind of show that no-one understood really how to use the padding_mask if it's just set to None in all the PRs. 

If we force ourselves to deprecate every internal function that doesn't have a `_` prefix (which many people still use in their codebases BTW) , it'll will be very hard to stay agile/dynamic. Ok for me to deprecate (just adding `kwargs` to all forward methods) even though I prefer not to and just state that we correct the `attention_mask` here.

@LysandreJik wdyt? 



Anything that isn't in the public API is private to us as maintainers, but after spending a bit of time looking in the documentation, I can't find where we have written it (but I remember us writing it down).

In any case the deprecation cycle doesn't cost much and I would keep it for 1-2 versions max before removing it. It was present for very few releases (one) and concerns very few users, so let's do a deprecation cycle with very explicit `v4.37.0` as a cutoff date.

This PR is ready for a review. All models that implemented FA2 are now refactored.

I've tested:
All slow tests for Falcon, Mistral, Llama they all pass on GPU.
```
CUDA_VISIBLE_DEVICES="0" RUN_SLOW=1 pytest tests/models/llama/test_modeling_llama.py
```

All flash attn tests:
```
CUDA_VISIBLE_DEVICES="0" RUN_SLOW=1 pytest -m flash_attn_test tests/models
```

There is one slow test that fails for Falcon, but it also fails on "main":
```
FAILED tests/models/falcon/test_modeling_falcon.py::FalconModelTest::test_left_padding_compatibility - AssertionError: False is not true
```
(might also be just my GPU)

As explained here: https://github.com/huggingface/transformers/pull/26792#discussion_r1366021714 this PR also fixes a bug with the sliding window mask for Mistral.


@ArthurZucker @younesbelkada @LysandreJik @fxmarty would be great if you could do a final review and run some tests to be sure everything works as expected! 

**Update**: I removed all the cache logic and instead just pass the attention_mask in the format that's needed. This is cleaner than caching tensors according to their shape, memory_id, etc...

All the benefits are kept including much improved readability and comprehensive attention mask class that can be copied / re-used by other models.

All tests pass just like before!

Then I am not sure what is the point of the class AttnMaskConverter? By the way, for SDPA, ideally we need both the information of 1/ is padding used 2/ transformers custom attention mask. This is because if custom masking is not used, we may dispatch on flash attention. So passing only a 4D mask for SDPA is suboptimal in my opinion.

Or I could just always pass the 4D attention mask to SDPA, but that kind of defeats the point given that dispatch to FA is then impossible.

> Are we adding the sliding window as a new feature for these models? Otherwise would just use two different classes for Mistral and the other

It's not possible really to use sliding window in Llama because it's hardcoded at initialization "sliding_window=...." for Mistral. So the user can't (and should not use) `sliding_window` for Llama via any config parameters (in the same way `is_causal` is hardcoded to `True`). It is true that that we only have a couple of architectures that use sliding windows (Mistral, Longformer, ...) so we could move it out of the attention mask converter class and instead put it directly into the forward for Mistral. I think it's better though to leave as is because:
- It's just a one-liner to add it to the attention converter and can be very nicely tested (which allowed us to spot the bug in Mistral)
- There is a high chance that we'll have more models with windowed attention if models build on Mistral
- We don't allow the user to configure windowed attention, so it's not like we adding a windowed attention feature to Llama or Falcon & thus making them more complicated.

But I do see how sliding window is arguably a bit exotic for the mask converter and if people feel strongly I can put it in Mistral's forward method instead.

Overall, we do move away a bit from "single-file" policy here as the attention converter is is a general class that has more features that needed for some models. But it does make sense here since there is really not much variation for attention mask across models and it greatly helps with readability.



No problem for me to leave the sliding window in the mask converter class, I indeed think we'll get to see more models leveraging the sliding window (or users that want it supported) in other architectures.