[FlexAttn] Fix models with unique characteristics

[vasqu]

For context, flex attention cannot work with dimensions less than 16; hence, the config was manipulated to ensure the test works. Before, most models failed, including llama.

There are some models such as idefics 2+3, smolvlm which do not have the _is_composite flag and as I do not want to affect other tests - so, I added a new condition to skip the test. They may have passed before but it's not future-proof. For Zamba2, I overwrote the test since some other dims don't add up when changing hidden_size.

There are other options:

• Rewrite the test to handle subconfigs --> tried that but there are so many edge cases and weird configs that lead to some issues one way or another.
• Adjust the dimensions in all models and avoid the hidden dim manipulation in the first place. Not sure if this is good as it will strain the tests even more imo 👀

Edit: #38434 took care of the composite models. This PR is left to fix some of the more unique models such as zamba2 and deepseek3.

[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.

[vasqu]

Maybe cc @zucchini-nlp (composite configs / models on vlms)

[vasqu]

#38434 seems to handle the vlms - will probably close this / only keep it for zamba2

[zucchini-nlp]

Good timing 😄

[ArthurZucker]

_is_composite is no longer needed / u8sed, but hasattr(config, "subconfigs") would be equivalent!

[vasqu]

I guess I unintentionally did that with len(config.sub_configs) > 1 lol

Should I merge this temporarily, or should we wait for #38434? Raushan's PR is trying to tackle the overall test to be composite compatible.

[ydshieh]

RUN_SLOW=1 python3 -m pytest -v tests/models/deepseek_v3/test_modeling_deepseek_v3.py::DeepseekV3ModelTest::test_flex_attention_with_grads

gives

ValueError: NYI: Currently non power of 2 embedding dimension are not supported. Got E=48 and Ev=32.

If you want to work on this too :-)

cc @zucchini-nlp too

[zucchini-nlp]

Btw @vasqu , my PR is merged so should be good to merge this one as well :)

[vasqu]

Nice, I'll take a look (maybe tomorrow, not sure) :D the deepseek issue looks like another hidden emb size issue 👀

Ig, I don't need to add the skip condition then?

[ydshieh]

no need to skip

[vasqu]

@ydshieh @zucchini-nlp If you want to take another look. The PR changed basically to fixing some more unique models which would be edge cases in the general test.

[zucchini-nlp]

LGTM, thanks