Persimmon FlashAttention2 [WIP]
Author: jeromeku
Pull request: yes

# What does this PR do?
Adds Flash Attention 2 for Persimmon per  #26350 

## Before submitting
- [x] This PR fixes a typo - see [commit](https://github.com/huggingface/transformers/commit/90d81d0379b9c6e0bc6643f2cd458fa5c7490fa3) 
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),

## Who can review?
@younesbelkada 

Ran tests on A100 80G, see attached for venv.
```
RUN_SLOW=1 pytest -sv --disable-warnings -k flash_attn_2 tests/models/persimmon/test_modeling_persimmon.py::PersimmonModelTest
================================================================================= test session starts ==================================================================================
platform linux -- Python 3.9.17, pytest-7.4.2, pluggy-1.3.0 -- /notebooks/virtualenvs/persimmon-fa2/bin/python
cachedir: .pytest_cache
configfile: setup.cfg
collecting ... Using /root/.cache/torch_extensions/py39_cu118 as PyTorch extensions root...
Detected CUDA files, patching ldflags
Emitting ninja build file /root/.cache/torch_extensions/py39_cu118/cuda_kernel/build.ninja...
Building extension module cuda_kernel...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module cuda_kernel...
Xformers is not installed correctly. If you want to use memory_efficient_attention to accelerate training use the following command to install Xformers
pip install xformers.
collected 116 items / 110 deselected / 6 selected                                                                                                                                      

tests/models/persimmon/test_modeling_persimmon.py::PersimmonModelTest::test_flash_attn_2_conversion You are attempting to use Flash Attention 2.0 with a model initialized on CPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
PASSED
tests/models/persimmon/test_modeling_persimmon.py::PersimmonModelTest::test_flash_attn_2_generate_left_padding You are attempting to use Flash Attention 2.0 with a model initialized on CPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
PASSED
tests/models/persimmon/test_modeling_persimmon.py::PersimmonModelTest::test_flash_attn_2_generate_padding_right You are attempting to use Flash Attention 2.0 with a model initialized on CPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
PASSED
tests/models/persimmon/test_modeling_persimmon.py::PersimmonModelTest::test_flash_attn_2_generate_use_cache You are attempting to use Flash Attention 2.0 with a model initialized on CPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
PASSED
tests/models/persimmon/test_modeling_persimmon.py::PersimmonModelTest::test_flash_attn_2_inference You are attempting to use Flash Attention 2.0 with a model initialized on CPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
You are attempting to use Flash Attention 2.0 with a model initialized on CPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
You are attempting to use Flash Attention 2.0 with a model initialized on CPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
PASSED
tests/models/persimmon/test_modeling_persimmon.py::PersimmonModelTest::test_flash_attn_2_inference_padding_right You are attempting to use Flash Attention 2.0 with a model initialized on CPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
You are attempting to use Flash Attention 2.0 with a model initialized on CPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
You are attempting to use Flash Attention 2.0 with a model initialized on CPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.
PASSED

==================================================================== 6 passed, 110 deselected, 8 warnings in 5.62s =====================================================================
```
[requirements.txt](https://github.com/huggingface/transformers/files/12761751/requirements.txt)
   

Comments:
@younesbelkada 
```
FAILED tests/models/persimmon/test_modeling_persimmon.py::PersimmonModelTest::test_pipeline_text_generation - AssertionError: (<class 'RuntimeError'>, <class 'IndexError'>, <class 'ValueError'>, <class 'AssertionError'>) not raised
```
This test  at `tests/pipelines/test_pipelines_text_generation.py:250: in run_pipeline_test
    text_generator("This is a test" * 500, max_new_tokens=20)` is failing since `Persimmon` should be able to handle long sequences.  The [default config](tests/models/persimmon/test_modeling_persimmon.py::PersimmonModelTest::test_pipeline_text_generation) has `max_position_embeddings = 16384` so the long sequence tests, which are run for `model_max_length < 1000` in `test_pipelines_text_generation` are not raising errors (so `assertRaises` is failing).

The other test failure is due to the tokenizer does not have a pad token:
```
FAILED tests/models/persimmon/test_modeling_persimmon.py::PersimmonModelTest::test_pipeline_zero_shot - ValueError: Asking to pad but the tokenizer does not have a padding token. Please select a token to use as `pad_token` `(tokenizer.pad_token = tokenizer.eos_token e.g.)` or add a new pad token via `tokenizer.add_special_tokens({'pad_token': '[PAD]'})`.
```
I'm guessing this is related to why the `PersimmonDecoderLayer` and `PersimmonFlashAttention` don't have a `kwarg` for `padding_mask` compared to the respective implementations for `LlamaDecoderLayer` and `LlamaFlashAttention`...

Thanks for your work on this! 
regarding the failing tests, this is very strange as the PR does not modify anything un-related to FA-2 so should keep all previous behaviour. Can you try to merge your branch with main and run the tests again?
Also can you elaborate on why pad/unpad is not needed for this architecture?

@younesbelkada 

Regarding the padding mask, both `LlamaDecoderLayer` and `LlamaFlashAttention` `forward` signatures have a `padding_mask` as a `kwarg` -- see [here](https://github.com/huggingface/transformers/blob/bd6205919aad4d3a2300a39a98a642f1cc3a5348/src/transformers/models/llama/modeling_llama.py#L614) and [here](https://github.com/huggingface/transformers/blob/bd6205919aad4d3a2300a39a98a642f1cc3a5348/src/transformers/models/llama/modeling_llama.py#L327).  

In contrast, neither the [`PersimmonDecoderLayer`](https://github.com/huggingface/transformers/blob/bd6205919aad4d3a2300a39a98a642f1cc3a5348/src/transformers/models/persimmon/modeling_persimmon.py#L371) nor the [`PersimmonFlashAttention`](https://github.com/huggingface/transformers/blob/bd6205919aad4d3a2300a39a98a642f1cc3a5348/src/transformers/models/persimmon/modeling_persimmon.py#L269) `forward` signatures have a `padding_mask` as a `kwarg`.

`LlamaFlashAttention2` per your implementation handles the cases with `padding` and `no padding` by calling two different methods of the `flash_attention_interface`, one which unpads, packs `qkv` and calls `flash_attn_varlen_func`), and in the other case with no padding `flash_attn_func`.

Since the `Persimmon` layers (per the original implementation) don't have a `padding_mask` as input, I only used `flash_attn_func` in the `PersimmonFlashAttention2` implementation.  Let me know if this needs to be changed.

@younesbelkada 
lmk if changes needed

@younesbelkada 

Crteated a new branch and PR #27052

Went ahead and added 2d-->4d attention mask per #26792 and adjusted FA2 to accommodate attention mask.

OK thanks! let's then close this PR in favor of #27052 ? what do you think?

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.

Any update on this?

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.