Support for SDPA for SAM models #34110

MagnusS0 · 2024-10-12T16:23:24Z

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.

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).
Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? 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, and
here are tips on formatting docstrings.
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.

HuggingFaceDocBuilderDev · 2024-10-22T14:24:58Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

ArthurZucker

Looks good! 🤗

src/transformers/models/sam/modeling_sam.py

ArthurZucker

Thanks for adding, I think this should be automatically tests, but make sure to run the slow tests

src/transformers/models/sam/configuration_sam.py

ArthurZucker · 2024-11-04T12:00:41Z

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!

MagnusS0 · 2024-11-04T12:53:01Z

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.

…_states

- Updated tests to skip cases flash and dynamic compile. - Minor adjustment to ensure correct loading of model with sdpa for dispatch test.

MagnusS0 · 2024-11-09T17:01:09Z

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.

qubvel · 2024-11-18T13:53:22Z

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?

qubvel · 2024-11-18T15:45:43Z

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

MagnusS0 · 2024-11-18T15:57:58Z

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?

qubvel · 2024-11-18T20:36:30Z

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

MagnusS0 · 2024-11-18T22:00:27Z

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

src/transformers/models/sam/configuration_sam.py

tests/models/sam/test_modeling_sam.py

…SAM model tests - Added attributes for sub_configs as per PR huggingface#34410. - Enabled tests for configs, ensuring the composite model (SAM) has several sub-configs in the main config. - Added class attribute _is_composite=True to the tester class - test_sdpa_can_dispatch_composite_models added

MagnusS0 · 2024-11-25T13:18:50Z

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

ArthurZucker

Thanks a lot for the addition, sorry for the delay.
We are going to refactor a bit afterwards in #35235, but let's merge for now! 🤗

ArthurZucker reviewed Oct 24, 2024

View reviewed changes

src/transformers/models/sam/modeling_sam.py Show resolved Hide resolved

ArthurZucker approved these changes Oct 31, 2024

View reviewed changes

src/transformers/models/sam/configuration_sam.py Outdated Show resolved Hide resolved

qubvel added Vision SDPA run-slow labels Oct 31, 2024

MagnusS0 force-pushed the sam_support_sdpa branch from 15192f3 to 5a16359 Compare November 9, 2024 15:54

MagnusS0 added 7 commits November 9, 2024 17:32

feat: add support for sdpa and gradient checkpointing

6dd779f

fix: ruff format

fcf3bcb

fix: config sdpa

40637ea

fix: sdpa layer naming convention

ee08617

fix: update test_eager_matches_sdpa_inference to handle vision_hidden…

e93fcfa

…_states

test: skip incompatible tests and fix loading issue with sdpa

70fe9bc

- Updated tests to skip cases flash and dynamic compile. - Minor adjustment to ensure correct loading of model with sdpa for dispatch test.

style: apply Ruff formatting

7ee64dc

MagnusS0 force-pushed the sam_support_sdpa branch from 8200114 to 7ee64dc Compare November 9, 2024 16:48

ruff fix again after rebase

df4e50c

Merge branch 'huggingface:main' into sam_support_sdpa

a2bff23

MagnusS0 requested a review from ArthurZucker November 14, 2024 21:34

MagnusS0 added 2 commits November 18, 2024 15:51

[run-slow] sam

c6ea29a

Merge branch 'main' into sam_support_sdpa

3a356b0

MagnusS0 added 2 commits November 18, 2024 16:58

Merge branch 'main' into sam_support_sdpa

0ed4b69

[run-slow] sam

69560e4

qubvel requested a review from zucchini-nlp November 18, 2024 20:36

zucchini-nlp reviewed Nov 19, 2024

View reviewed changes

src/transformers/models/sam/configuration_sam.py Outdated Show resolved Hide resolved

tests/models/sam/test_modeling_sam.py Outdated Show resolved Hide resolved

MagnusS0 added 6 commits November 20, 2024 21:29

[run-slow] sam

b401e94

style: ruff

b470746

[run-slow] sam

e12c188

style: ruff again ...

dbd416c

[run-slow] sam

d572d19

ArthurZucker reviewed Dec 17, 2024

View reviewed changes

ArthurZucker merged commit 6eb00dd into huggingface:main Dec 17, 2024
27 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support for SDPA for SAM models #34110

Support for SDPA for SAM models #34110

MagnusS0 commented Oct 12, 2024 •

edited

Loading

HuggingFaceDocBuilderDev commented Oct 22, 2024

ArthurZucker left a comment

ArthurZucker left a comment

ArthurZucker commented Nov 4, 2024

MagnusS0 commented Nov 4, 2024 •

edited

Loading

MagnusS0 commented Nov 9, 2024

qubvel commented Nov 18, 2024

qubvel commented Nov 18, 2024 •

edited

Loading

MagnusS0 commented Nov 18, 2024 •

edited

Loading

qubvel commented Nov 18, 2024

MagnusS0 commented Nov 18, 2024

MagnusS0 commented Nov 25, 2024

ArthurZucker left a comment

Support for SDPA for SAM models #34110

Support for SDPA for SAM models #34110

Conversation

MagnusS0 commented Oct 12, 2024 • edited Loading

What does this PR do?

Performance Improvements (Updated 03.11.2024)

Before submitting

Who can review?

Some TODO

HuggingFaceDocBuilderDev commented Oct 22, 2024

ArthurZucker left a comment

Choose a reason for hiding this comment

ArthurZucker left a comment

Choose a reason for hiding this comment

ArthurZucker commented Nov 4, 2024

MagnusS0 commented Nov 4, 2024 • edited Loading

MagnusS0 commented Nov 9, 2024

qubvel commented Nov 18, 2024

qubvel commented Nov 18, 2024 • edited Loading

MagnusS0 commented Nov 18, 2024 • edited Loading

qubvel commented Nov 18, 2024

MagnusS0 commented Nov 18, 2024

MagnusS0 commented Nov 25, 2024

ArthurZucker left a comment

Choose a reason for hiding this comment

MagnusS0 commented Oct 12, 2024 •

edited

Loading

MagnusS0 commented Nov 4, 2024 •

edited

Loading

qubvel commented Nov 18, 2024 •

edited

Loading

MagnusS0 commented Nov 18, 2024 •

edited

Loading