Update tests regarding attention types after #35235

[ydshieh]

What does this PR do?

As discussed offline, let's update test_flash_attn_2_can_dispatch_composite_models after #35235.

I keep the original condition and using or (....new condition), although the new condition should work well anyway.

[zucchini-nlp]

Thanks! I think we can entirely delete that because we already checked the config in lines 4441-4444

And since we started fixing these, can we also remove the check on class name from similar tests with something like "SdpaAttention" in name" ?

[ydshieh]

ok i will give it a shot

[HuggingFaceDocBuilderDev]

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.

[ydshieh]

@zucchini-nlp The check (a few lines above) only check model_fa2.config (not the submodules, and in particular, not the attention modules), but for the most strict check, it might be worth to check the "Attention" layer module have this config attribute and config.flash_attention_2 has the same value we expect.

Do you still think it's better to remove the part you suggested?

[zucchini-nlp]

The above check should verify that all sub-model have fa2, and would be enough if we are sure the attn implementation is propagated further correctly

it might be worth to check the "Attention" layer module have this config attribute

I guess this one makes more sense for new attention, so yeah we can also check attention module's config and the attribute value there

[ydshieh]

(implicitly assuming such attention module always have config attribute, otherwise it might suggest something is not propagating correctly to attn module/layer)

[ydshieh]

let's me know if the latest version is what you expect :-)

[ydshieh]

if so I will apply similar changes to other places as you suggested

[zucchini-nlp]

Yep, LGTM as long as the tests are green :)

[ydshieh]

I updated other tests. It turns out that

(implicitly assuming such attention module always have config attribute, otherwise it might suggest something is not propagating correctly to attn module/layer)

will fail for some classes, see below. I will use getattr(..., None) to pass those cases.

Failing test.

FAILED tests/models/xlm_roberta_xl/test_modeling_xlm_roberta_xl.py::XLMRobertaXLModelTest::test_sdpa_can_dispatch_non_composite_models - AttributeError: 'XLMRobertaXLAttention' object has no attribute 'config'
FAILED tests/models/bert/test_modeling_bert.py::BertModelTest::test_sdpa_can_dispatch_non_composite_models - AttributeError: 'BertAttention' object has no attribute 'config'
FAILED tests/models/kosmos2/test_modeling_kosmos2.py::Kosmos2ModelTest::test_attn_implementation_composite_models - AttributeError: 'KosmosTextAttention' object has no attribute 'config'
FAILED tests/models/kosmos2/test_modeling_kosmos2.py::Kosmos2ModelTest::test_sdpa_can_dispatch_composite_models - AttributeError: 'KosmosTextAttention' object has no attribute 'config'
FAILED tests/models/vivit/test_modeling_vivit.py::VivitModelTest::test_sdpa_can_dispatch_non_composite_models - AttributeError: 'VivitAttention' object has no attribute 'config'
FAILED tests/models/vit_mae/test_modeling_vit_mae.py::ViTMAEModelTest::test_sdpa_can_dispatch_non_composite_models - AttributeError: 'ViTMAEAttention' object has no attribute 'config'
FAILED tests/models/videomae/test_modeling_videomae.py::VideoMAEModelTest::test_sdpa_can_dispatch_non_composite_models - AttributeError: 'VideoMAEAttention' object has no attribute 'config'
FAILED tests/models/dinov2_with_registers/test_modeling_dinov2_with_registers.py::Dinov2WithRegistersModelTest::test_sdpa_can_dispatch_non_composite_models - AttributeError: 'Dinov2WithRegistersAttention' object has no attribute 'config'
FAILED tests/models/yolos/test_modeling_yolos.py::YolosModelTest::test_sdpa_can_dispatch_non_composite_models - AttributeError: 'YolosAttention' object has no attribute 'config'
FAILED tests/models/beit/test_modeling_beit.py::BeitModelTest::test_sdpa_can_dispatch_non_composite_models - AttributeError: 'BeitAttention' object has no attribute 'config'

[zucchini-nlp]

LGTM! Btw, are the fa2 tests passing for all models? IMO we need to add the getattr() hack in fa2 test as well

[ydshieh]

LGTM! Btw, are the fa2 tests passing for all models? IMO we need to add the getattr() hack in fa2 test as well

Good catch, you are right. I was a bit lazy and after seeing CircleCI is green I didn't think in more depth. Guess CircleCI didn't run that test FA2 test and I should have run it manually.

Once verified all passing, I will merge!

[ydshieh]

There is a single failing

tests/models/sam/test_modeling_sam.py::SamModelTest::test_flash_attn_2_can_dispatch_composite_models

that is already failing on main.

I will merge and come back to this SAM model separately later.