Allow passing 2D attention mask
Author: UniverseFly
Pull request: no
Labels: name: Feature request

### Feature request

Allow passing a 2D attention mask in `model.forward`.

### Motivation

With this feature, it would be much easier to avoid cross-context contamination during pretraining and supervised finetuning when packing the sequences together for more efficient training.

Here is an example usecase discussed in (https://github.com/huggingface/trl/issues/805):

![](https://user-images.githubusercontent.com/26831266/272305004-93c690a8-7e9b-40ad-885f-d530996aa109.png)

### Your contribution

Upon investigation into the source code, I found the current logic of initializing attention masks is mostly a fixed code snippet encoded in each model:

```python
        if getattr(self.config, "_flash_attn_2_enabled", False):
            # 2d mask is passed through the layers
            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
        else:
            # 4d mask is passed through the layers
            attention_mask = _prepare_4d_causal_attention_mask(
                attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
            )
```

To enable this behavior may require hacking into each model. I should be able to handle part of them and submit a draft PR. But before that, I want to know if this feature request is reasonable.

Comments:
Hey, the model's forward already supports passing a 2d attention mask, it is just expended to 4d because that is the format required by the attention implementation. 
Would you mind elaborating on what you cannot currently do? (Might be related to #27539?) 

> Hey, the model's forward already supports passing a 2d attention mask, it is just expended to 4d because that is the format required by the attention implementation.
> Would you mind elaborating on what you cannot currently do? (Might be related to #27539?)

Yeah, I might not make it clear. The current "2D"s are `[batch_size, num_tokens]`. What I suggested was `[batch_size, num_tokens, num_tokens]` so we can have a matrix for each batch that explicitly defines what each token should attend to. https://github.com/huggingface/transformers/pull/27539 seems relevant

Just chiming in, here is some more context (also very interested in this feature). From what I understand, this is not trivial implement in general...

As one current example, the [axolotl](https://github.com/OpenAccess-AI-Collective/axolotl/tree/main) finetuning harness implements efficient sample packing with correct block diagonal attention masking through [a series of monkey patches](https://github.com/OpenAccess-AI-Collective/axolotl/tree/main/src/axolotl/monkeypatch) for the underlying huggingface model definitions for a few of the very popular models like llama and mistral. Though I have not looked through the code in detail, I believe it leverages the fact that the flash attention api supports the masking required to implement this scheme.

It is relevant for efficient finetuning (the reason it's incorporated into axolotl), and general wisdom (and whispers from inside large corps) suggest that this type of block diagonal masking is better for large scale training code.

(https://github.com/huggingface/transformers/pull/27539 is relevant, but it looks like the focus may be on the beam search/speculative decoding use case, not this slightly more general use case. Also here's a relevant hf forum post https://discuss.huggingface.co/t/the-correct-attention-mask-for-examples-packing/52909/2)

Packing is indeed a good use-case for supporting 2D attention mask for huggingface models.

Packing is planned

> Packing is planned

Hello, is there any detailed schedule to support this feature ? many thanks.

Most probably not next release, but the one after that!

Looking forward to this feature!

#31446 for packing 

Hi @ArthurZucker, does #31446 include packing? It seems it is just refactoring flash attention, a prerequisite of packing not packing itself.

Yep, it's planned not done yet. I was gonna do both but ended up splitting!

cc @Cyrilvallez 

#33932 is related for the packing as well! 

Any progress on this?

Hey! We do not yet support generation/training with packed sequences, however in https://github.com/huggingface/transformers/pull/35235 we applied the necessary changes to most frequent model architectures (and more are coming). Thus, as of now you could feed ragged tensors (just add a dummy dimension of one in front of the 1d input tensor) directly to the model initialized with FA2, along with the seq_lens args of FA2 as kwargs. 
 
Otherwise, you can use sdpa with 4d attention mask of shapes [batch, 1, seq_len, seq_len] to achieve what you want 🤗

We are soon going to start using packed sequences natively whenever FA2 is enabled 🤗

Hi @Cyrilvallez,
Thank you for your input! I noticed there’s already a blog post about using packing with FA2: <https://huggingface.co/blog/packing-with-FA2>. Does this mean the feature is already supported?

I’m relatively new to deep learning, so apologies if this is a silly question: Is Flash Attention required to use packing? I currently only have access to Volta GPUs, which are not supported by Flash Attention. Is it possible to implement packing without relying on Flash Attention?


Hey! For now out of all the attention functions we have available in `Transformers`, only flash attention supports packing natively. However, you could define your own `flex_attention` implementation of packing! Check out https://pytorch.org/blog/flexattention/ for some examples/help on how to do it. Then you can use your custom attention function directly with your `Transformers` model by just registering it 🤗