using custom attention mask

[whatisslove11]

Hello!
I am doing a translation task and would like to try using flash attention in my model
In addition to the usual triangular mask, I also need to mask padding tokens so that the model does not pay attention to them - sequences of the same length already arrive in the model itself
As I understand it, there is no function of feeding your mask yet
Could you tell me how I can use my mask or make flash attention add padding tokens to the mask itself?

[whatisslove11]

@tridao

[tridao]

You can look at how we do it in BERT: Remove all padding tokens before the first layer.

Idk if that works for translation.

[flower-with-safe]

what if I want to use reset-attention-mask when pretrain a llama model?
for example, my attention mask could be:
tensor([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[1., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
[1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],
[1., 1., 1., 1., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 1., 1., 0., 0., 0., 0.],
[0., 0., 0., 0., 1., 1., 1., 0., 0., 0.],
[0., 0., 0., 0., 1., 1., 1., 1., 0., 0.],
[0., 0., 0., 0., 1., 1., 1., 1., 1., 0.],
[0., 0., 0., 0., 1., 1., 1., 1., 1., 1.]])
in such case how could I use flash attention?
@tridao

[sentialx]

@tridao Is it possible to use a PrefixLM attention mask?

[tridao]

No that's not supported.

[sentialx]

It looks like custom attention mask is the most requested feature. In fact, it seems to be already solved by this PR: #57

[iiLaurens]

Am also much interested in custom masks. I think the value of Prefixlm mask is not appreciated enough. Would like to experiment continuation of pretraining with prefixlm a la UL2RUL2R.

@tridao, are you considering support for a custom attention masks? Or do you have specific objections to it?