Support for SDPA for SAM models
Author: MagnusS0
Pull request: yes
Labels: name: Vision

name: SDPA

name: run-slow

# What does this PR do?

This PR adds support for SDPA to the SAM (Segment Anything Model), following the implementation done by the PyTorch team in [segment-anything-fast](https://github.com/pytorch-labs/segment-anything-fast/).

I add two new classes:

- `SamSdpaAttention(SamAttention)`: For the decoder.
- `SamSdpaVisionAttention(SamVisionAttention)`: For the encoder.

### Performance Improvements (Updated 03.11.2024)
Time Reduction between Eager and SDPA modes for SAM-B and SlimSAM averaged across 10 runs in ms:

| **Batch Size** | **Eager (ms)**      | **SDPA (ms)**       | **Time Reduction (%)** |
|----------------|---------------------|---------------------|------------------------|
| **SAM-B**      |                     |                     |                        |
| 1              | 64.58 (±2.14)       | 30.95 (±0.07)       | 52.06%                 |
| 8              | 1,021.45 (±4.55)    | 227.44 (±0.67)      | 77.73%                 |
| **SlimSAM**    |                     |                     |                        |
| 1              | 50.51 (±1.19)       | 28.70 (±0.03)       | 43.16%                 |
| 8              | 913.97 (±6.73)      | 284.43 (±2.14)      | 68.87%                 |


Peak Memory Usage During Forward Pass:

| **Batch Size** | **Eager (MB)** | **SDPA (MB)** | **Memory Reduction (%)** |
|----------------|----------------|---------------|--------------------------|
| **SAM-B**      |                |               |                          |
| 1              | 2,151          | 875           | 59.31%                   |
| 8              | 16,007         | 3,913         | 75.54%                   |
| **SlimSAM**    |                |               |                          |
| 1              | 1,961          | 1,273         | 35.10%                   |
| 8              | 15,540         | 9,503         | 38.84%                   |


On the TrashNet 1.0 validation set I got a reduction around 30-35%, with no decrease in accuracy.
I also ran training on TrashNet 1.0 with no issues and about 1.7x speedup per epoch with SlimSAM. Also allowed gradient check-pointing that seemed to be missing from the model deceleration, worked without any changes.

## 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?
- [ ] 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.
- [ ] 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).
- [ ] Did you write any new necessary tests?


## Who can review?

@amyeroberts 
@qubvel 
I'll also tag @ArthurZucker since we just worked on this model last weekend.

### Some TODO
Should probably add some documentation for this, but let me know where is best place to add this. 

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

```FAILED tests/models/sam/test_modeling_sam.py::SamModelTest::test_sdpa_can_dispatch_non_composite_models - ValueError: The SDPA model should have SDPA attention layers``` can't merge until it's fixed!

> `FAILED tests/models/sam/test_modeling_sam.py::SamModelTest::test_sdpa_can_dispatch_non_composite_models - ValueError: The SDPA model should have SDPA attention layers` can't merge until it's fixed!

Doing run slow I also get the following failed tests:
```
FAILED tests/models/sam/test_modeling_sam.py::SamModelTest::test_eager_matches_sdpa_inference_0_float16 - AttributeError: 'SamImageSegmentationOutput' object has no attribute 'hidde...
FAILED tests/models/sam/test_modeling_sam.py::SamModelTest::test_eager_matches_sdpa_inference_1_bfloat16 - AttributeError: 'SamImageSegmentationOutput' object has no attribute 'hidde...
FAILED tests/models/sam/test_modeling_sam.py::SamModelTest::test_eager_matches_sdpa_inference_2_float32 - AttributeError: 'SamImageSegmentationOutput' object has no attribute 'hidde...
FAILED tests/models/sam/test_modeling_sam.py::SamModelTest::test_sdpa_can_compile_dynamic - torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
FAILED tests/models/sam/test_modeling_sam.py::SamModelTest::test_sdpa_can_dispatch_non_composite_models - ValueError: The SDPA model should have SDPA attention layers
FAILED tests/models/sam/test_modeling_sam.py::SamModelTest::test_sdpa_can_dispatch_on_flash - RuntimeError: No available kernel. Aborting execution.
FAILED tests/models/sam/test_modeling_sam.py::SamModelIntegrationTest::test_inference_mask_generation_batched_points_batched_images - AssertionError: False is not true
FAILED tests/models/sam/test_modeling_sam.py::SamModelIntegrationTest::test_inference_mask_generation_no_point - AssertionError: False is not true
FAILED tests/models/sam/test_modeling_sam.py::SamModelIntegrationTest::test_inference_mask_generation_one_point_one_bb - AssertionError: False is not true
FAILED tests/models/sam/test_modeling_sam.py::SamModelIntegrationTest::test_inference_mask_generation_one_point_one_bb_zero - AssertionError: False is not true
```

- The 4 tests related to `test_inference_mask_generation` fail on the main branch as well. 
- I think `test_sdpa_can_dispatch_on_flash` should be skipped since the flash kernel can't be run with attention_mask's that is used to add the relative positional embeddings. 
- `test_eager_matches_sdpa_inference` is missing hidden states, can the vision_hidden_states be used for SAM in this test?

I haven't found the root cause of `test_sdpa_can_dispatch_non_composite_models`, running it outside the test framework it does not fail so I believe it might be related to how 'sdpa' is added to the configs (??)


I'll work on a fix for those I think I'll be able to solve.


Apologies for the commit history mess. I had to rebase to clean things up, so the branch should now be up to date with the main branch and tests pass.

Hi @MagnusS0, thanks for the update, can you please push an empty commit with the message `[run-slow] sam` to trigger slow tests in CI?

Sorry, it should be the last commit, in that case I can approve the run 🙂

> Sorry, it should be the last commit, in that case I can approve the run 🙂

Ahh no problem, will commit again.  Just fyi there are 4 slow tests I expect will fail, as they also do that on main. 
They will pass if the `assertion` `atol=1e-4` is set to `atol=1e-3`. Would this be an fair fix?


It looks like everything is fine in CI 🤗 Let's wait for a comment from @zucchini-nlp regarding setting attention implementation in config and I suppose it's ready to be merged

Awesome 🤗  Probably just some small differences in my local setup causing the test to fail then.

@ArthurZucker Could you please review this PR when you have a moment? It should be ready to merge 🤗 Thank you!