add sdpa to OPT
Author: avishaiElmakies
Pull request: yes
Labels: name: run-slow

adds SDPA to OPT model

impl inspired by ```gemma2``` and ```llama```.

part of #28005.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?
 I think @amyeroberts @fxmarty 

some notes:

- I did some refactoring to the model code. created a function ```_update_key_and_values```. the code was used by all 3 attention implementations. ```self._shape``` to ```_shape```. and moved masks logic to ```self._update_casual_mask```.
- I created a test that makes sure the generate of eager and SDPA are equivalent. the test is similar to the one ```llama```
- I seem to fail 3 implementations of common tests ```test_eager_matches_sdpa_inference_0_float16```, ```test_eager_matches_sdpa_inference_1_bfloat16```, ```test_eager_matches_sdpa_inference_2_float32```. I took inspiration from gemma, which seems to ignore those tests as well. should i ignore them as well? I think it is related also related to #32086 since my code is similar to that. Also seems the affect equivalence with flax/tf beacuse the default will become sdpa. 

would love some feedback!

Comments:
@amyeroberts fixed the comments. would love to understand what to do with the tests that are failing

@avishaiElmakies Similar to [my comment](https://github.com/huggingface/transformers/pull/33403#issuecomment-2340236947) for DinoV2 -- other PRs can be a good reference for how to fix. There you will see that many of the tests enforce the test model to use eager attention. This will resolve the TF/flax equivalence issues. 

For the sdpa equivalence tests, I'm not sure. As the flash attention implementation follows llama, following llama's tests for sdpa seems reasonable cc @ArthurZucker here who might now more about this 

Hi @amyeroberts. looked at the faiiling test

changed the test ```test_eager_matches_sdpa_generate``` to match the one in ```llama```, fixed ```test_pt_tf_model_equivalence```.

the tests in the file ```test_modelling_tf_opt.py```. ```test_xla_generate_contrastive``` and ```test_xla_generate_slow```. don't seem related to this PR. they also fail in main.

the only tests which I would love some guidance on are the ```test_eager_matches_sdpa_inference```. generating seems to work the same with both implementations. from what I understand and after looking at it the behavior for both implementations is different in the inference test. as stated by #32086. fixing this might require a refactor, to make sure behaviors match.

would love some guidance. 

EDIT: I noticed when running the tests using your setup, it fails on ```test_eager_matches_sdpa_generate```. but when running the test on my machine it passes. 
![image](https://github.com/user-attachments/assets/fed31e90-9b11-4191-95b6-0a31ebf972e5)


The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_33298). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

The code looks eerily similar to bart's implementation, it might be more beneficial to let it "inherit" from bart via `# Copied from ...` statements and adjust it to fit in with barts uses of attention masks etc instead of manually figuring stuff out 👀 

Hence, it's also more appropriate to look into encoder-decoder models like bart instead of llama who are decoder only regarding the attention implementation.

Ok #17437 seems to be the cause why the copied from doesn't exist in the first place. Is there a way to just ignore a few lines since it seems a bit too much to me to remove the whole copied from then 👀 

@vasqu I can take inspiration from Bart instead of llama,if that helps. What do you think? 

@avishaiElmakies I think we can closely follow bart, maybe even leverage a `Copied from` for the SDPA attention module. It also makes sense to follow bart regarding the attention mask creation then. It's more tested and makes our life easier :)

For the generation, they can be flaky so I'd opt (unintended word play) for up to a 1b ish model for the test. 125M is still relatively small for todays standards. 

@vasqu OK, I will look into changing it more to a Bart inspiration.

I will also try to use use a larger model, i think opt has 125M, 350M and 1.3B
Do you want the 1.3B? I also hope I have the resources to test with the larger models 😅

Hopefully, the 350M suffices no need to go further than 1.3b tho. For the larger models, we can leverage the slow runs if that's necessary ;) 

@vasqu OK, thanks for the guidance! I will work on this and update when I'm Done

@vasqu I changed the test to 350M.

I also think i found the bug thanks to bart, the 5 tests are now passing (at least on my machine)

@avishaiElmakies great to hear! Let's wait until Amy has time to review but I'll take a closer look in a bit :) 

Thx LGTM, the failure on CI is not related to this PR.

@amyeroberts would love a review, a slow-run test and hopefully a merge 🫡

@amyeroberts is there some code to run those benchmarks?

EDIT: nvm think i found it.

@amyeroberts  ran benchmarks

used code from here https://github.com/huggingface/transformers/pull/31031

local resources - (L40S-45GB, PyTorch 2.4.0, OS Debian GNU/Linux 11

### Training 

|    batch_size |    seq_len |  Time per batch (eager - s)   |    Time per batch (sdpa - s) |  Speedup (%)   |  Eager peak mem (MB)   |    sdpa peak mem (MB) |  Mem saving (%)   |
|--------------:|-----------:|:------------------------------|-----------------------------:|:---------------|:-----------------------|----------------------:|:------------------|
|             1 |        128 | 0.047                         |                        0.037 | 26.360         | 1474.611               |               1474.32 | 0.019             |
|             1 |        256 | 0.046                         |                        0.037 | 24.335         | 1498.541               |               1499.49 | -0.063            |
|             1 |        512 | 0.046                         |                        0.037 | 24.959         | 1973.544               |               1551.35 | 27.215            |
|             1 |       1024 | 0.062                         |                        0.038 | 65.135         | 4867.113               |               1698.35 | 186.578           |
|             1 |       2048 | 0.230                         |                        0.039 | 483.933        | 15662.224              |               2715.75 | 476.718           |
|             2 |        128 | 0.045                         |                        0.037 | 20.455         | 1498.164               |               1499.49 | -0.089            |
|             2 |        256 | 0.046                         |                        0.037 | 24.027         | 1569.367               |               1551.35 | 1.161             |
|             2 |        512 | 0.045                         |                        0.037 | 20.965         | 3257.074               |               1698.35 | 91.778            |
|             2 |       1024 | 0.122                         |                        0.038 | 225.958        | 9054.405               |               2715.75 | 233.403           |
|             2 |       2048 | 0.464                         |                        0.067 | 593.646        | 30572.058              |               4750.55 | 543.548           |
|             4 |        128 | 0.045                         |                        0.037 | 21.918         | 1549.448               |               1551.35 | -0.123            |
|             4 |        256 | 0.044                         |                        0.038 | 18.084         | 2451.768               |               1698.35 | 44.361            |
|             4 |        512 | 0.069                         |                        0.037 | 84.421         | 5833.180               |               2715.75 | 114.791           |
|             4 |       1024 | 0.262                         |                        0.062 | 319.475        | 17427.842              |               4750.55 | 266.860           |
|             4 |       2048 | OOM                           |                        0.062 | Eager OOM      | OOM                    |               4750.55 | Eager OOM         |
|             8 |        128 | 0.044                         |                        0.037 | 18.436         | 2049.115               |               1697.78 | 20.694            |
|             8 |        256 | 0.048                         |                        0.036 | 32.887         | 4222.567               |               2715.75 | 55.484            |
|             8 |        512 | 0.153                         |                        0.06  | 154.862        | 10985.391              |               4750.55 | 131.245           |
|             8 |       1024 | 0.526                         |                        0.122 | 330.697        | 34175.763              |               8821.18 | 287.428           |
|             8 |       2048 | OOM                           |                        0.122 | Eager OOM      | OOM                    |               8821.18 | Eager OOM         |

### Inference

|    batch_size |    seq_len |    Per token latency eager (ms) |    Per token latency SDPA (ms) |    Speedup (%) |    Mem eager (MB) |    Mem BT (MB) |    Mem saved (%) |
|--------------:|-----------:|--------------------------------:|-------------------------------:|---------------:|------------------:|---------------:|-----------------:|
|             1 |        128 |                          11.634 |                          8.647 |         34.546 |           717.676 |        717.674 |            0     |
|             1 |        256 |                          11.593 |                          8.86  |         30.851 |           742.852 |        742.845 |            0.001 |
|             1 |        512 |                          11.515 |                          8.816 |         30.614 |           798.232 |        799.593 |           -0.17  |
|             1 |       1024 |                          11.556 |                          8.915 |         29.628 |           917.265 |        895.538 |            2.426 |
|             2 |        128 |                          12.724 |                         11.002 |         15.659 |           762.434 |        762.431 |            0     |
|             2 |        256 |                          12.704 |                         11.063 |         14.83  |           816.809 |        816.733 |            0.009 |
|             2 |        512 |                          12.757 |                         10.947 |         16.535 |           917.383 |        918.339 |           -0.104 |
|             2 |       1024 |                          13.018 |                         11.018 |         18.147 |          1162.65  |       1114.81  |            4.291 |
|             4 |        128 |                          12.739 |                         10.959 |         16.243 |           856.335 |        856.483 |           -0.017 |
|             4 |        256 |                          12.718 |                         10.837 |         17.355 |           957.298 |        957.674 |           -0.039 |
|             4 |        512 |                          12.813 |                         10.822 |         18.393 |          1158.44  |       1158.45  |           -0.001 |
|             4 |       1024 |                          13.416 |                         11.06  |         21.301 |          1653.42  |       1557.19  |            6.18  |
|             8 |        128 |                          12.763 |                         10.891 |         17.193 |          1036.13  |       1036.51  |           -0.036 |
|             8 |        256 |                          12.89  |                         11.104 |         16.085 |          1236.98  |       1236.87  |            0.01  |
|             8 |        512 |                          13.327 |                         10.939 |         21.836 |          1642.29  |       1641.78  |            0.031 |
|             8 |       1024 |                          15.181 |                         11.175 |         35.848 |          2634.98  |       2443.35  |            7.843 |


updated model card opt.md as well



@avishaiElmakies Can you make the `[slow-run] opt` commit; it's needed to start slow runs. Great benchmarks btw!

Ci errors are not related.

@vasqu commited.

Yeah, the benchmarks look great. I was not expecting this much of an improvement when training 

slow runs cc @amyeroberts 

@amyeroberts it seems it have skipped the tests, do you know what happend? 

@avishaiElmakies I think it's because the message in the commit message is `[slow-run]` rather than `[run-slow]`. Could you try again with `[run-slow] opt`? 

@amyeroberts sorry about that, commited with `[run-slow] opt`

The failures seem unrelated to me tbh. XLA should be tf stuff, no? @amyeroberts 

@vasqu Yes, it should be TF stuff, although the TF files are touched in this PR so they should be passing too. Could you confirm if the tests are passing on main or not? 

cc @Rocketknight1 @ArthurZucker 

@amyeroberts at the time those tests were failing in main as well

Confirmed that `test_xla_generate_contrastive` and `test_xla_generate_slow` are failing on `main` for me in TF - I think it's fine to just add skips to those tests in this PR for now, since it's not related to this PR.

Passing locally on my machine with a GPU ([tf_opt_logs.txt](https://github.com/user-attachments/files/17232562/tf_opt_logs.txt)) and not with a CPU ([tf_opt_logs_cpu.txt](https://github.com/user-attachments/files/17232603/tf_opt_logs_cpu.txt)).

Seems like a CPU issue then? FYI @Rocketknight1 
cc @amyeroberts 

Yep, we can add the wrapper `require_accelerator`! 

@ArthurZucker There's a proper fix in #33903 since the GPU is just silently ignoring errors instead of throwing errors in TF even though we go beyond the possible positional embedding range.

@ArthurZucker @amyeroberts do i need to do anything or can this be merged? 

Sorry for the delay and thanks for the contribution! 🤗