4D `attention_mask` support
Author: poedator
Pull request: yes

This is implementation for feature request from #27493 [custom 4d attention_mask as transformers .forward() argument](https://github.com/huggingface/transformers/issues/27493).

1) Allowing 4d attention masks to pass thru `_prepare_4d_causal_attention_mask()` intact
2) support in OPT (need to build custom `positions` tensor)
3) support in Llama (while Llama can accept custom `position_ids`, I added code to generate them internally)

The benefits of the code are to enable more memory-efficient text generation with tree-based parallel decoding as described in [SpecInfer paper](https://arxiv.org/abs/2305.09781)

Tagging:
@gante (generate)
@patrickvonplaten (masks)
@younesbelkada @ArthurZucker (text models)

This PR is WiP:
- Will add tests
- Need advice on how to handle models beyond covered Llama and OPT
- May add example for memory-efficient generation

**IMPORTANT: this PR makes changes that can only used by few classes of models**
requirements to use:
- have `position_ids` argument in `.forward()` method
- use `modeling_attn_mask_utils.py::_prepare_4d_attention_mask()` function for 4d mask generation
- 

as of 20.12.2023, only a handful (under 20) of transformers model classes meet these criteria. Most of these classes are multimodal, which may require their own use cases for 4D masks. The pure language modelling classes fit to use the 4D mask changes from this PR are only `LlamaModel`, `FalconModel` and `XGLMModel`.

Comments:
Generally, I don't have a problem with allowing to pass 4D attention masks! @poedator can you explain your use case a little bit for why you want to pass 4d attention masks? 

@patrickvonplaten 
here is a use example:
Suppose one does beam search and has a starting prefix with tokens `11 22 33` in 4 beams. Now he needs to check candidates with tokens 44, 55, 66, and 77. Present code would pack the beams into a batch of shape (4, 4):

```
11 22 33  44
11 22 33  55
11 22 33  66
11 22 33  77
```
and run it with mask of all ones, passing such mask in 2D which gets expanded internally to 4D.

The proposed way would be to have a batch shaped (1, 7):
`11 22 33  44 55 66 77`
and the 4d mask would have a shape (1, 1, 4, 7) and look like this:
```
1  1  1  1  0  0  0 
1  1  1  0  1  0  0 
1  1  1  0  0  1  0 
1  1  1  0  0  0  1

with a positions tensor of [0, 1, 2, 3, 3, 3, 3]
```

At subsequent beam search iterations the mask will reflect which past tokens should the new tokens attend to.
Such mask needs to pass intact. 
This saves memory for past_key_values cache and thus allows beam search and other similar inference (like SpecInfer) of longer sequences with limited VRAM.

Another use case is kindly proposed by @UniverseFly below.

Very interesting PR! Would this feature also enable SFT packing as mentioned in https://github.com/huggingface/trl/issues/805?


![](https://user-images.githubusercontent.com/26831266/272305004-93c690a8-7e9b-40ad-885f-d530996aa109.png)


> Very interesting PR! Would this feature also enable SFT packing as mentioned in [huggingface/trl#805](https://github.com/huggingface/trl/issues/805)?
Sure it would. Just have a separate packing function somewhere - it is beyond the scope of this PR. 
Besides, one should be able to pack multiple series of sequences into a batch this way. 


I tried this branch and the `model.forward` seems to work fairly well, but `model.generate` raises errors with the 4D attention mask (with Llama). After some checking, it might be due to the missing logic here:

https://github.com/huggingface/transformers/blob/53a7e7750ff088ffbd7d96c5aeed122cc96b6866/src/transformers/models/llama/modeling_llama.py#L1087-L1124

Generate looks like a harder challenge for your methods - each individual sequence will be expanding, thus you'd need to reorder past_kv and mask at each step. I believe that to implement it, you'd need to write custom `prepare_inputs_for_generation()`, and possibly some more logic. 
I'll be happy to test drive it. 
On my side I intend to write a PR for more efficient beam search after this PR merges.

Hi, @ArthurZucker  
I limited this PR only to the mask code, proceeding with the tests. 

So far I have demo in [Colab with monkey patch based on this PR.](https://colab.research.google.com/drive/1PMcLNjjjK0Zwgg6rQJdutwolfIVCZfs6?usp=sharing) It shows a negligible difference in logits obtained the old and new ways. I dent to believe that this is a rounding error somewhere. Would you support it as the basis for the tests? 
BTW, where to put this new test?

Hi, @UniverseFly ,
Try the monkey patch from the [Colab notebook](https://colab.research.google.com/drive/1PMcLNjjjK0Zwgg6rQJdutwolfIVCZfs6?usp=sharing) - see if it works to implement your idea.  

Thanks for this PR and the demo. It is very helpful in trying the [SpecInfer paper](https://arxiv.org/abs/2305.09781). Also in another recent progress on speculative decoding [look ahead decoding](https://lmsys.org/blog/2023-11-21-lookahead-decoding/) Fig 5, this PR will also be useful.

Reviewing now 😉 

- squashed all earlier commits into one
- added tests. Made a separate class to test with full model loading.
- added support for sdpa (following https://github.com/huggingface/transformers/pull/26572)
- `test_modeling_utils.py::AttentionMaskTester` and `::TestAttentionImplementation` tests pass
- new tests pass

@ArthurZucker, please review. Hopefully it is ready to merge.

@ArthurZucker, pls give me a hint about `NameError: name 'torch' is not defined` error. Apparently a decorator or import is missing, but can't figure it out. The import and decorators seem in place...

> Hi, @ArthurZucker I limited this PR only to the mask code, proceeding with the tests.
> 
> So far I have demo in [Colab with monkey patch based on this PR.](https://colab.research.google.com/drive/1PMcLNjjjK0Zwgg6rQJdutwolfIVCZfs6?usp=sharing) It shows a negligible difference in logits obtained the old and new ways. I dent to believe that this is a rounding error somewhere. Would you support it as the basis for the tests? BTW, where to put this new test?
> 
> Hi, @UniverseFly , Try the monkey patch from the [Colab notebook](https://colab.research.google.com/drive/1PMcLNjjjK0Zwgg6rQJdutwolfIVCZfs6?usp=sharing) - see if it works to implement your idea.

@poedator, Did you ever get the 4D beam search working with the monkey patch in the Colab notebook? I would be very interested if you were able to get this working already! (especially depth first kv cache updating)

> @poedator, Did you ever get the 4D beam search working with the monkey patch in the Colab notebook? I would be very interested if you were able to get this working already! (especially depth first kv cache updating)
Hi, @Codys12! Thank you for the interest to the PR.
I made a working implementation of memory-efficient beam search using this 4D mask. Got as high as 256 beams of 32 tokens  with Llama 7 on A100 and this is far from the limit. There is no complete demo to share, but [this gist ](https://gist.github.com/poedator/c754247f3dca8f70b710186c9bc37032)has the beam search part of my code. Hope that you find it useful.


having a look!

> Sorry checking the test there are duplicate markers 😅 not sure they are needed no?

Earlier, I got frustrated with failing commits and added decorators everywhere. Now most of them are gone and it still passes CI checks.

Thanks for the contribution! 🤗 

@ArthurZucker , would you want to publish a blog post in HF blog with 4d attention use cases?
I propose to include:
- memory efficient beam search (my example, from tests)
- SFT packing as mentioned in https://github.com/huggingface/trl/issues/805, suggested by @UniverseFly 
- [look ahead decoding](https://lmsys.org/blog/2023-11-21-lookahead-decoding/), suggested by @KexinFeng 


If you want feel free to do so! 🤗 

Note that not all paths of this can be `torch.compile`d:

The following fails due to `torch.all(attention_mask == 1)`.

```
import torch
import torch.nn as nn
from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask_for_sdpa

class Model(nn.Module):
    def forward(self, inputs_embeds):
        batch_size, seq_length, _ = inputs_embeds.shape
        past_key_values_length = 10
        attention_mask = torch.tensor([1.])
        attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
            attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
        )
        return attention_mask

model = Model()
model = torch.compile(model, fullgraph=True)
model(torch.ones([1,5, 32]))
```

@PhilJd,
`torch.all(attention_mask == 1)` was present even before this PR.
[see this line](https://github.com/huggingface/transformers/pull/27539/files#diff-b14be70e49c04e876d7cd745948bf1bee279bc0d6f2a71b1a18e3b5aff293bd1L343)
it comes form https://github.com/huggingface/transformers/pull/26572

have you tested the preceding commit?

Ah sorry, just looked at the blame - yeah, the previous commit fails as well @fxmarty .

`_prepare_4d_causal_attention_mask` is applied only if `self._use_flash_attention_2` is False (https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L1039). Is it because 4D attention mask and flash attention 2 are not compatible?

The function description should be updated to avoid confusion as `attention_mask` is not necessarily 2D now (https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_attn_mask_utils.py#L290)

@shentianxiao , thank you for your attention to the 4D attention!

> `_prepare_4d_causal_attention_mask` is applied only if `self._use_flash_attention_2` is False (https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L1039). Is it because 4D attention mask and flash attention 2 are not compatible?

it is not about compatibility, rather the flash_attention_2 code contrasted original mask vs modified mask coming from `_prepare_4d_causal_attention_mask()`

> The function description should be updated to avoid confusion as attention_mask is not necessarily 2D now (https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_attn_mask_utils.py#L290)

I agree, that the original mask may also be 4d-shaped now. I just started PR https://github.com/huggingface/transformers/pull/28151 with documentation updates - will make edits there. Hopefully the maintainers responsible for `flash_attention_2` will verify it.

**IMPORTANT: this PR makes changes that can only used by few classes of models**
requirements to use:
- have `position_ids` argument in `.forward()` method
- use `modeling_attn_mask_utils.py::_prepare_4d_attention_mask()` function for 4d mask generation
- 

as of 20.12.2023, only a handful (under 20) of transformers model classes meet these criteria. Most of these classes are multimodal, which may require their own use cases for 4D masks. The pure language modelling classes fit to use the 4D mask changes from this PR are only `LlamaModel`, `FalconModel` and `XGLMModel`.

I made a small blog post based on this PR. 
https://huggingface.co/blog/poedator/4d-masks
Big thanks to everyone who contributed and commented!

Thanks for the amazing addition!! This is a great new feature.

Just wanted to ask a question to make sure I am using it properly. In the code [here](https://github.com/poedator/transformers/blob/d80cb9823cc8b774bb4f41ac59579edca8f79ff0/src/transformers/modeling_attn_mask_utils.py#L357), it looks like the 4D masks are expected to have shape `[batch_size, 1, seq_len, seq_len].` (I am inferring that the `1` in the `expected_shape` is the heads dimension so that the same mask is broadcast to all heads.) In the [blog post](https://huggingface.co/blog/poedator/4d-masks), it describes the attention masks as having shape `[heads, batch_size, input_ids_length, total_sequence_length]`.

My question is: **are the `heads` and `batch_size` dimensions transposed in the blog post**? It seems like we are actually expected to provide 4D masks where the first axis is batch size, the second is heads. The blog post implies the reverse. Since I am sometimes using a batch size of 1 in testing, this works either way, but I want to use it correctly and don't see the "proper" shape documented anywhere (perhaps it is documented somewhere and I missed it!).

Thanks!

@jpgard ,
you are correct, there was an error in my blog post.
Changed it to `[batch_size, heads, input_ids_length, total_sequence_length]`
thank you for raising this!

Great, thanks for the quick reply and for your hard work on this @poedator !!

Has this been tested with flash attention 2? Works great for me without flash attention 2, but when using flash attention I get lots of messages of the form `../aten/src/ATen/native/cuda/ScatterGatherKernel.cu:144: operator(): block: [202,0,0], thread: [105,0,0] Assertion idx_dim >= 0 && idx_dim < index_size && "index out of bounds" failed.`

Lower chunk of the stack trace posted below.

```
 File "/admin/home-jpgard/miniconda3/envs/rtfm/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 798, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/admin/home-jpgard/miniconda3/envs/rtfm/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/admin/home-jpgard/miniconda3/envs/rtfm/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/admin/home-jpgard/miniconda3/envs/rtfm/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 549, in forward
    attn_output = self._flash_attention_forward(
  File "/admin/home-jpgard/miniconda3/envs/rtfm/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 592, in _flash_attention_forward
    query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
  File "/admin/home-jpgard/miniconda3/envs/rtfm/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 631, in _upad_input
    query_layer = index_first_axis(
  File "/admin/home-jpgard/miniconda3/envs/rtfm/lib/python3.8/site-packages/torch/autograd/function.py", line 553, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/admin/home-jpgard/miniconda3/envs/rtfm/lib/python3.8/site-packages/flash_attn/bert_padding.py", line 17, in forward
    return torch.gather(
RuntimeError: CUDA error: device-side assert triggered
```

Would be great to be able to use FA2 with this PR as the speedups are much larger as sequence length grows -- so FA2 seems like the perfect accompaniment to e.g. "packed" training sequences enabled by this PR.

@jpgard , please share some simple testing code. I will look into this issue.



Hello! I have a few questions about using custom 4D attention masks for sequence packing:

Does the custom 4D attention mask need to explicitly encode causality (e.g., enforcing a lower-triangular structure for autoregressive tasks)?
For the custom 4D attention mask, is the correct format (batch_size, 1, seq_len, seq_len), where valid attention positions are marked with 1, and invalid ones with 0?
Additionally, I noticed that enabling sequence packing with a custom 4D attention tensor results in a lower average loss compared to the non-packed scenario, even when I explicitly enforce lower-diagonality in the mask. Could you clarify why this might be the case?

Thanks in advance for your help!

Hi, @DylanZSZ ,

> Does the custom 4D attention mask need to explicitly encode causality
yes, it should have the lower triangular (or more complex for special cases) structure for causality. 

>where valid attention positions are marked with 1, and invalid ones with 0? 
actually use 0 for valid and -inf for invalid. This is how the masks are applied in transformers after some internal transformations. 

> enabling sequence packing with a custom 4D attention tensor results in a lower average loss compared to the non-packed scenario, even when I explicitly enforce lower-diagonality in the mask. Could you clarify why this might be the case?
Hard to tell, it should be same in theory.  suggestions: (1) make sure you use 0 and -inf, as above, (2) try with some simpler examples and follow the attention in regular case and with custom mask. (3) check if padding is done properly.