Llama: fix custom 4D masks

[gante]

What does this PR do?

Fixes the issue raised by @poedator in this comment.

Causal mask is now of shape [..., seq_len, full_len], as opposed to [..., full_len, full_len]. This means custom 4D attention masks are now the whole causal mask, so we don't need a sliced copy -- we can copy the whole thing :)

This PR also expands the support of custom 4D attention mask: we can pass both the full mask ([..., full_len, full_len]) or the partial mask ([..., seq_len, full_len]).

[gante]

reordered the logic: custom 4D masks are now a superset of the default mask, so we don't need to create the default mask first :)

[gante]

This else has no changes. Only the if attention_mask is not None and attention_mask.dim() == 4: is different.

[gante]

Added this test case (we can now pass full custom 4D attention masks)

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

[ArthurZucker]

Thanks LGTM just want to always trigger the tests

[poedator]

Thanks @gante !
I tested the code and it works.
My only suggestion is to cover the 4D masks in the documentation for Llama and other models.

[gante]

My only suggestion is to cover the 4D masks in the documentation for Llama and other models.

@poedator would you like to open a PR with that? As a user, you'll probably have cool examples in mind!

[poedator]

My only suggestion is to cover the 4D masks in the documentation for Llama and other models.

@poedator would you like to open a PR with that? As a user, you'll probably have cool examples in mind!

will try, but not this week...

[gante]

This set of slow tests was moved to the llama test file -> if we run the slow llama tests, which we often request, this will now be triggered

[gante]

This set of tests are now:

1. part of the mixin, so they are run on all push commits
2. a fast test, using model = model_class(config) from the test config
3. triggered by model_class._supports_cache_class == True -- recent LLMs [llama, cohere, gemma, mistral, mixtral, starcoder2, ...] have this attribute set to True and are 4D mask-compatible. Older models are often not compatible. Over time, as we spread the cache refactor, this test will be run on those classes as well 👀

[gante]

@ArthurZucker ready for a re-review (test rework) -- we now have on push tests for all recent models + custom 4D mask :)

[poedator]

@gante ,
May I suggest adding extra tests for 4D masks with StaticCache? I am concerned that StaticCache code may handle the custom masks differently. I intend to use 4D masks with StaticCache in my new speculative decoding implementation.
Here are the additional methods for Mask4DTestHard:
https://gist.github.com/poedator/f1c15551d202df2682c65c1bbdcb1c07

I made the cache longer than the masks and padded the masks to the cache length. is this the correct way?

[ArthurZucker]

Sorry for the delay, let's rebase on main as well

[ArthurZucker]

Very good! Let's rebase on main, #30047 was merged, and run slow tests!

[poedator]

please, please merge this PR - I need it for my speculative decoding paper project! The 4D masks are essential for it.
@gante @ArthurZucker

[ArthurZucker]

Sorry just got back to github 😓 could you rebase!

[poedator]

Sorry just got back to github 😓 could you rebase!
Too bad it did not make it into 4.40 (

I rebased this PR in new one #30348 and added few important changes.

[poedator]

@gante, thank you for the rebase. Meanwhile I added more improvements in #30348 - let's close this #29930 and continue there.

[gante]

Closing in favor of #30348