transformers
transformers copied to clipboard
Llama: fix custom 4D masks, v2
this is an attempt to rebase #29930 started initially by @gante
Fixes the issue raised by @poedator in https://github.com/huggingface/transformers/pull/29753#issuecomment-2014530814.
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]).
as of 18.04.24 it is not passing the 4D mask tests because of the _ignore_causal_mask_sdpa()
method, most recently edited in #30317 (merged today).
tests/models/llama/test_modeling_llama.py::Mask4DTestHard::test_partial_stacked_causal_mask - ValueError: Incorrect 4D attention_mask shape: (1, 1, 12, 12); expected: (1, 1, 9, 12).
apparently _ignore_causal_mask_sdpa()
expects that attention_mask.shape[-2] == query_length
. This may only be true if the input_ids are contiguous, which is not always the case in some intended 4D mask applications.
tests/models/llama/test_modeling_llama.py::Mask4DTestHard::test_stacked_causal_mask_static_cache - ValueError: Incorrect 4D attention_mask shape: (1, 1, 12, 16); expected: (1, 1, 12, 12).
in this test _ignore_causal_mask_sdpa()
expects that attention_mask.shape[-1] == key_value_length
which is set by past seen tokens
. However, in the test I make this dimension equal to the static cache size, so that the mask always has same shape and the whole graph may be compiled and reused.
I hesitate to make edits to _ignore_causal_mask_sdpa()
because there may be some greater context.
Summoning @younesbelkada @ArthurZucker @gante to help
As s a solution, I added additional expected_shapes
to _ignore_causal_mask_sdpa()
and improved StaticCache detection code.
Note: it is inconvenient to have StaticCache as layer.self_attn objects and other Caches as model-level object. Perhaps there may be a model-level plug to avoid referencing the layer levels.
Please review soon - I need this for my paper code. It's been broken for quite long now.
The LONG tests look OK.
all CI tests are green, SLOW tests were OK on my side yesterday
I noticed that mistral model support for 4D masks stayed broken after these fixes. So I added similar lines to src/transformers/modeling_attn_mask_utils.py::_prepare_4d_causal_attention_mask_for_sdpa()
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.
I added Mask4DTestHard
tests (without static cache part) to tests/models/mistral/test_modeling_mistral.py
to ensure that the 4d masks keep working in the models that use _prepare_4d_causal_attention_mask_for_sdpa()
. These new tests would fail without the fixes from commit d488f35 just above.
Tested the SLOW tests for ./tests/models/mistral/
branch - all fine
Is there anything left to do before the merge is possible? @gante @ArthurZucker
Let's remove unrelated changes!
sorry, but without these changes, the fixes and tests will not work. I looked for related PRs, all I found was #30476 but it is not fixing the relevant parts of the code.
I tried to follow Arthur's advice to streamline the path for the 4D masks and it seems to work. The relevant tests do pass. @ArthurZucker @gante , please review
I combined the 2 tests from common, which were very similar. Added tolerance - now Mixtral passes it OK. @ArthurZucker, @gante - please see if it is good to merge now
Just waiting for the commits to be push to check the done
parts and we can merge