F.scaled_dot_product_attention support
Author: fxmarty
Pull request: yes

As per title, this PR proposes to support natively `torch.nn.functional.scaled_dot_product_attention` in transformers. I propose to enable SDPA by default if `torch>=2.1.1` (released 15 Nov. 2023), for the reasons written in the PR. The support could then be extended using https://github.com/huggingface/optimum/blob/main/optimum/bettertransformer/models/attention.py.

---

The introduced `_unmask_unattended` is a workaround for https://github.com/pytorch/pytorch/issues/110213.

It behaves as follow:

If attention_mask is
 ```
    [[0, 0, 1]
     [1, 1, 1]
     [0, 1, 1]]
```
and expanded_mask is (e.g. here left-padding case)
```
    [[[[0, 0, 0],
       [0, 0, 0],
       [0, 0, 1]]],
     [[[1, 0, 0],
       [1, 1, 0],
       [1, 1, 1]]],
     [[[0, 0, 0],
       [0, 1, 0],
       [0, 1, 1]]]]
```
then the modified expanded_mask will be
```
    [[[[1, 1, 1],   <-- modified
       [1, 1, 1],   <-- modified
       [0, 0, 1]]],
     [[[1, 0, 0],
       [1, 1, 0],
       [1, 1, 1]]],
     [[[1, 1, 1],   <-- modified
       [0, 1, 0],
       [0, 1, 1]]]]
```

Modifying as such the attention mask is fine given that we modify it only for pad tokens on the `-2` dimension. Softmax is computed on the `-1` dimension, and thus there is no change for the relevant non-padding tokens.

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_26572). All of your documentation changes will be reflected on that endpoint.

As https://github.com/huggingface/transformers/pull/26792 was merged will get back to it this week, targeting next to next transformers release.

I think it would be wise to put a requirement on `torch>=2.1.1` due to this issue https://github.com/pytorch/pytorch/issues/112577, but you are the judge. Happy to do a check on torch version and call `contiguous()` of we want to support `2.1.0`.

This PR adds the support of SDPA for the following architectures, activated by default with `torch>=2.1.1`:
* Llama
* Falcon (support extended to alibi, bugfix with attention row of mask all masked)
* GPT BigCode
* Bart
* Whisper

The other model files changed are the result of  `Copied from`. None of them are removed.

The method `_unmask_unattended` added to `AttentionMaskConverter` is critical to avoid this bug https://github.com/pytorch/pytorch/issues/110213.

The addition of the methods `_prepare_4d_causal_attention_mask_for_sdpa` and `_prepare_4d_attention_mask_for_sdpa` are useful to dispatch to FA1/FA2 that can not handle a non `None` `attn_mask` argument.

The `"default"` in `ATTENTION_CLASSES` is replaced by `"eager"`, as the default will be dependent on PyTorch version. Maybe there is a better name to find than `"eager"`.

I argue that this issue https://github.com/pytorch/pytorch/issues/112577 is a good enough argument to put a requirement on `torch>2.1.1` (not yet released).

There are likely still bugs - but I think it is a good time for review @patrickvonplaten @younesbelkada @amyeroberts @ArthurZucker @LysandreJik 

@ArthurZucker @patrickvonplaten @LysandreJik @amyeroberts feel free to rereview, this should be in good shape

@ArthurZucker thanks a lot for the great points raised in your review. They should be addressed.

@patrickvonplaten @LysandreJik happy to hear your opinion as well

Before merge (or whenever it's the time to check), please ping me to trigger CI on GPU. Thanks.

For sure @ydshieh!

> Looks good to me, but tend to think that the mask transformation might be better in the AttentionClass rather that in the model class (as we have if else for each implementation basically). It's probably the last blocker for me with regard to the modeling codes.

Yes that's a great point, and related to some discussions in https://github.com/huggingface/transformers/pull/26792 & https://huggingface.slack.com/archives/C060RADBR4J/p1697050725287089 (private, you should be in). My understanding is that @patrickvonplaten at first wanted to do exactly that, moving the attention_mask logic to `LlamaAttention` class, with a caching mechanism to avoid recomputing at each layer: https://github.com/huggingface/transformers/pull/26792#issuecomment-1771836936

However, my understanding is the caching mechanism was not very elegant, which just motivated the if/else in a forward at the LlamaModel level. I agree it is not super super ideal. @patrickvonplaten what was the issue with https://github.com/huggingface/transformers/blob/ae3eb2e72aa0c0e8fad8df66e0b4199168c98619/src/transformers/models/llama/modeling_llama.py#L93? That `tensor.data_ptr()` is not enough and we need actually a `torch.equal()` to check that the mask is the same?

What is the final API now that the user can use to enable SDPA? 

@patrickvonplaten 
> I propose to enable SDPA by default if torch>=2.1.1 (released 15 Nov. 2023), for the reasons written in the PR.

Apart from that, exposed API is `attn_implementation="sdpa"`, `attn_implementation="eager"` that may be used e.g. to disable SDPA if required (head_mask, output_attentions).

We may need to agree on the name `attn_implementation` and `eager`.

`config.attn_implementation` is made private to `config._attn_implementation` in https://github.com/huggingface/transformers/pull/26572/commits/5c77b944749268b11f30c01a010d84fd692d670a as suggested

`torch.jit.trace` is unhappy about our implementation https://github.com/pytorch/pytorch/issues/115262, making some tests fail. It's a tricky one.

@fxmarty can we maybe just disable `torch.jit.trace` for testing now? Don't think it's super important tbh

@patrickvonplaten Indeed unfortunately that's what I go for for now: one need `attn_implementation="eager"` (or `torch<2.1.1`) to use `torch.jit.trace` successfully for architectures that support SDPA.

This PR should be in good shape.

Left to do:
- [x] Falcon with torch==2.0 needs to use SDPA anymore by default for BC
- [x] Make `torch.jit.trace` work with SDPA attention when an attention mask is provided. Tracing SDPA with torch.jit.trace does not work by default with models supporting SDPA and loaded with `torch>=2.1.1` (see https://github.com/pytorch/pytorch/issues/115262) when no attention_mask is passed. Models should be loaded with `attn_implementation="eager"` to be exported with `torch.jit.trace`, or use `attention_mask` when tracing. This may be improved in the future.

For the record, the `is_causal` controlflow is needed due to https://github.com/pytorch/pytorch/issues/108108 & the fact that we want as much as possible be able to dispatch to flash attention (passing `attn_mask` makes it impossible).

As SDPA is still going through changes in PyTorch (see https://github.com/pytorch/pytorch/pull/114823 & https://github.com/pytorch/pytorch/issues/110681), I think we should keep the flexibility of the minimum `torch` version required to use those classes (e.g. we may want to bump to `torch>=2.2` in the future.

It is ready. Here is a summary of the relevant CI.

#### `RUN_SLOW=1 CUDA_VISIBLE_DEVICES=0 pytest tests/models/llama -s -vvvvv`

Flacky (new):
```
FAILED tests/models/llama/test_modeling_llama.py::LlamaModelTest::test_eager_matches_sdpa_inference_1_bfloat16 - AssertionError: False is not true : padding_side=left, use_mask=False, batch_size=5, enable_kernels=True: mean relative difference: 8.728e-03, torch ato...
```

Already failing on `main`:
```
FAILED tests/models/llama/test_modeling_llama.py::CodeLlamaIntegrationTest::test_model_7b_logits - AssertionError: Lists differ: ['
▁<PRE> def remove_non_ascii(s: str) -> st[893 chars]ID>'] != ['
 <PRE> def remove_non_ascii(s: str) -> st[893 chars...
FAILED tests/models/llama/test_tokenization_llama.py::LlamaIntegrationTest::test_conversion - AssertionError: '{\n [964 chars]or": {\n    "type": "TemplateProcessing",\n   [1795198 chars]}\n}' != '{\n [964 chars]or": null,\n  "decoder": {\n    "t...
```

#### `RUN_SLOW=1 CUDA_VISIBLE_DEVICES=0 pytest tests/models/whisper -s -vvvvv`

Already failing on `main`:
```
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperModelIntegrationTests::test_large_generation_multilingual - FileNotFoundError: https://voice-prod-bundler-ee1969a6ce8178826482b88e843c335139bd3fb4.s3.amazonaws.com/cv-corpus-6.1-2020-12-11/ja.tar.gz
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperModelIntegrationTests::test_whisper_longform_multi_batch - assert [' While Porashaggy sits there, a cooing dove. He has gone, and gone for good," answered Polychrom, who had managed to squeeze into the room besi...
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperModelIntegrationTests::test_whisper_longform_multi_batch_hard - assert " Folks, if you watch the show, you know, I spent a lot of time right over there. Patiently and astutely scrutinizing the boxwood and mahogany ch...
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperModelIntegrationTests::test_whisper_longform_single_batch - assert [" Because you were sleeping instead of conquering, the lovely rose princess has become a fiddle without a bow, all poor ashaggy sits there, acco...
```

Flacky (on `main`):
```
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperStandaloneDecoderModelTest::test_flash_attn_2_generate_left_padding - AssertionError: False is not true
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperStandaloneDecoderModelTest::test_flash_attn_2_inference - AssertionError: assert False
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperStandaloneDecoderModelTest::test_flash_attn_2_inference_padding_right - AssertionError: assert False
```

#### `RUN_SLOW=1 CUDA_VISIBLE_DEVICES=0 pytest tests/models/bart -s -vvvvv`

Flacky on `main`:
```
AILED tests/models/bart/test_modeling_bart.py::BartStandaloneDecoderModelTest::test_cpu_offload - AssertionError: False is not true
```

#### `RUN_SLOW=1 CUDA_VISIBLE_DEVICES=0 pytest tests/models/falcon -s -vvvvv`

Flacky (new):
```
FAILED tests/models/falcon/test_modeling_falcon.py::FalconModelTest::test_eager_matches_sdpa_inference_1_bfloat16 - AssertionError: False is not true : padding_side=left, use_mask=True, batch_size=5, enable_kernels=True: mean relative difference: 7.141e-03, torch atol...
```

Flacky (on `main`):
```
FAILED tests/models/falcon/test_modeling_falcon.py::FalconModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
```

#### `RUN_SLOW=1 CUDA_VISIBLE_DEVICES=0 pytest tests/models/idefics -s -vvvvv`

all pass

#### `RUN_SLOW=1 CUDA_VISIBLE_DEVICES=0 pytest tests/models/bert -s -vvvvv`

all pass

#### `RUN_SLOW=1 CUDA_VISIBLE_DEVICES=0 pytest tests/models/gpt2 -s -vvvvv`

all pass

#### `RUN_SLOW=1 CUDA_VISIBLE_DEVICES=0 pytest tests/test_modeling_utils.py -s -vvvvv`

Already failing on `main`:
```
FAILED tests/test_modeling_utils.py::ModelUtilsTest::test_legacy_load_from_url - huggingface_hub.utils._validators.HFValidationError: Repo id must be in the form 'repo_name' or 'namespace/repo_name': 'https://huggingface.co/hf-intern...
FAILED tests/test_modeling_utils.py::ModelUtilsTest::test_load_from_one_file - huggingface_hub.utils._validators.HFValidationError: Repo id must be in the form 'repo_name' or 'namespace/repo_name': '/tmp/tmp64wrpwyf'. Use `repo_typ...
FAILED tests/test_modeling_utils.py::ModelUtilsTest::test_model_from_pretrained - AssertionError: 7 != 8
ERROR tests/test_modeling_utils.py::ModelOnTheFlyConversionTester::test_safetensors_on_the_fly_conversion - ValueError: Cannot run tests as secret isn't setup.
ERROR tests/test_modeling_utils.py::ModelOnTheFlyConversionTester::test_safetensors_on_the_fly_conversion_gated - ValueError: Cannot run tests as secret isn't setup.
ERROR tests/test_modeling_utils.py::ModelOnTheFlyConversionTester::test_safetensors_on_the_fly_conversion_private - ValueError: Cannot run tests as secret isn't setup.
ERROR tests/test_modeling_utils.py::ModelOnTheFlyConversionTester::test_safetensors_on_the_fly_sharded_conversion - ValueError: Cannot run tests as secret isn't setup.
ERROR tests/test_modeling_utils.py::ModelOnTheFlyConversionTester::test_safetensors_on_the_fly_sharded_conversion_gated - ValueError: Cannot run tests as secret isn't setup.
ERROR tests/test_modeling_utils.py::ModelOnTheFlyConversionTester::test_safetensors_on_the_fly_sharded_conversion_private - ValueError: Cannot run tests as secret isn't setup.
ERROR tests/test_modeling_utils.py::ModelOnTheFlyConversionTester::test_safetensors_on_the_fly_specific_revision - ValueError: Cannot run tests as secret isn't setup.
ERROR tests/test_modeling_utils.py::ModelOnTheFlyConversionTester::test_safetensors_on_the_fly_wrong_user_opened_pr - ValueError: Cannot run tests as secret isn't setup.
```

#### `RUN_SLOW=1 CUDA_VISIBLE_DEVICES=0 pytest tests/ -s -vvvvv -k "flash or sdpa"`

Flacky (new):
```
FAILED tests/models/falcon/test_modeling_falcon.py::FalconModelTest::test_eager_matches_sdpa_inference_1_bfloat16 - AssertionError: False is not true : padding_side=left, use_mask=False, batch_size=1, enable_kernels=True: mean relative difference: 7.660e-03, torch ato...
```

Already failing/flacky on `main`:
```
FAILED tests/models/bark/test_modeling_bark.py::BarkSemanticModelTest::test_flash_attn_2_from_config - ValueError: Unrecognized configuration class <class 'transformers.models.bark.configuration_bark.BarkSemanticConfig'> for this kind of AutoModel: AutoMo...
FAILED tests/models/bark/test_modeling_bark.py::BarkSemanticModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
FAILED tests/models/bark/test_modeling_bark.py::BarkCoarseModelTest::test_flash_attn_2_from_config - ValueError: Unrecognized configuration class <class 'transformers.models.bark.configuration_bark.BarkCoarseConfig'> for this kind of AutoModel: AutoMode...
FAILED tests/models/distilbert/test_modeling_distilbert.py::DistilBertModelTest::test_flash_attn_2_inference_padding_right - AssertionError: False is not true
FAILED tests/models/gpt_bigcode/test_modeling_gpt_bigcode.py::GPTBigCodeModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
FAILED tests/models/gpt_neo/test_modeling_gpt_neo.py::GPTNeoModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
FAILED tests/models/gpt_neox/test_modeling_gpt_neox.py::GPTNeoXModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
FAILED tests/models/opt/test_modeling_opt.py::OPTModelTest::test_flash_attn_2_generate_padding_right - AssertionError: False is not true
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperStandaloneDecoderModelTest::test_flash_attn_2_inference - AssertionError: assert False
FAILED tests/models/whisper/test_modeling_whisper.py::WhisperStandaloneDecoderModelTest::test_flash_attn_2_inference_padding_right - AssertionError: assert False
```